Update leveldb

Signed-off-by: Yuan Si <do4suki@gmail.com>

database/leveldb.go

@@ -1,68 +0,0 @@
-package database
-
-import (
-	"errors"
-
-	"github.com/syndtr/goleveldb/leveldb"
-	"github.com/syndtr/goleveldb/leveldb/util"
-)
-
-type LevelDbImpl struct {
-	ldb *leveldb.DB
-}
-
-var (
-	LevelDb LevelDbImpl
-	err     error
-)
-
-func init() {
-	LevelDb.InitDb()
-}
-
-func (ldb *LevelDbImpl) InitDb() {
-	ldb.ldb, err = leveldb.OpenFile("./data/honoka-chan.db", nil)
-	if err != nil {
-		panic(err.Error())
-	}
-}
-
-func (ldb *LevelDbImpl) Get(key []byte) (res []byte, err error) {
-	if len(key) == 0 {
-		return nil, errors.New("key is null")
-	}
-	res, err = ldb.ldb.Get(key, nil)
-	if err != nil {
-		return nil, err
-	}
-	return res, nil
-}
-
-func (ldb *LevelDbImpl) Put(key, value []byte) (err error) {
-	if len(key) == 0 {
-		return errors.New("key is null")
-	}
-	err = ldb.ldb.Put(key, value, nil)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-func (ldb *LevelDbImpl) List() (res map[string]string) {
-	res = make(map[string]string)
-	iter := ldb.ldb.NewIterator(nil, nil)
-	for iter.Next() {
-		res[string(iter.Key())] = string(iter.Value())
-	}
-	return res
-}
-
-func (ldb *LevelDbImpl) ListPrefix(prefix []byte) (res map[string]string) {
-	res = make(map[string]string)
-	iter := ldb.ldb.NewIterator(util.BytesPrefix(prefix), nil)
-	for iter.Next() {
-		res[string(iter.Key())] = string(iter.Value())
-	}
-	return res
-}

database/utils.go

@@ -1,18 +0,0 @@
-package database
-
-import "fmt"
-
-func MatchTokenUid(token, uid string) bool {
-	// ret := LevelDb.List()
-	// for k, v := range ret {
-	// 	fmt.Println(k, v)
-	// 	fmt.Println()
-	// }
-	res, err := LevelDb.Get([]byte(uid))
-	if err != nil {
-		fmt.Println(err)
-		return false
-	}
-
-	return string(res) == token
-}

db/db.go

@@ -0,0 +1,21 @@
+package db
+
+import "fmt"
+
+var (
+	DB *Instance
+)
+
+func init() {
+	DB = GetInstance()
+}
+
+func MatchTokenUid(token, uid string) bool {
+	res, err := DB.Get([]byte(uid))
+	if err != nil {
+		fmt.Println(err)
+		return false
+	}
+
+	return string(res) == token
+}

db/leveldb.go

@@ -0,0 +1,58 @@
+package db
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/syndtr/goleveldb/leveldb"
+)
+
+var (
+	once sync.Once
+)
+
+type Instance struct {
+	db *leveldb.DB
+	mu sync.Mutex
+}
+
+func GetInstance() *Instance {
+	in := Instance{}
+	once.Do(func() {
+		var err error
+		in.db, err = leveldb.OpenFile("./data/honoka-chan.db", nil)
+		if err != nil {
+			panic(err)
+		}
+	})
+	return &in
+}
+
+func (in *Instance) Close() error {
+	in.mu.Lock()
+	defer in.mu.Unlock()
+	if in.db != nil {
+		return in.db.Close()
+	}
+	return nil
+}
+
+func (in *Instance) Get(key []byte) ([]byte, error) {
+	if len(key) == 0 {
+		return nil, errors.New("empty key")
+	}
+	res, err := in.db.Get(key, nil)
+	if err == leveldb.ErrNotFound {
+		return nil, nil
+	} else if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (in *Instance) Set(key, value []byte) error {
+	if len(key) == 0 {
+		return errors.New("empty key")
+	}
+	return in.db.Put(key, value, nil)
+}

handler/live.go

@@ -3,7 +3,7 @@ package handler
 import (
 	"encoding/json"
 	"honoka-chan/config"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"honoka-chan/model"
 	"honoka-chan/utils"
 	"math"
@@ -34,7 +34,7 @@ func PlayLive(ctx *gin.Context) {
 
 	// Save Deck Id for /live/reward
 	key := "live_deck_" + ctx.GetString("userid")
-	err = database.LevelDb.Put([]byte(key), []byte(strconv.Itoa(deckId)))
+	err = db.DB.Set([]byte(key), []byte(strconv.Itoa(deckId)))
 	CheckErr(err)
 
 	// Song type: normal / special
@@ -525,7 +525,7 @@ func PlayReward(ctx *gin.Context) {
 	CheckErr(err)
 
 	key := "live_deck_" + ctx.GetString("userid")
-	deckId, err := database.LevelDb.Get([]byte(key))
+	deckId, err := db.DB.Get([]byte(key))
 	CheckErr(err)
 	unitsList := []model.PlayRewardUnitList{}
 	err = UserEng.Table("deck_unit_m").Join("LEFT", "user_deck_m", "deck_unit_m.user_deck_id = user_deck_m.id").

handler/login.go

@@ -2,7 +2,7 @@ package handler
 
 import (
 	"encoding/json"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"honoka-chan/model"
 	"net/http"
 	"strconv"
@@ -38,7 +38,7 @@ func Login(ctx *gin.Context) {
 	}
 	ctx.Set("userid", userId)
 
-	err = database.LevelDb.Put([]byte(strconv.Itoa(userId)), []byte(ctx.GetString("authorize_token")))
+	err = db.DB.Set([]byte(strconv.Itoa(userId)), []byte(ctx.GetString("authorize_token")))
 	CheckErr(err)
 
 	loginResp := model.LoginResp{

handler/private.go

@@ -5,7 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"honoka-chan/config"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"honoka-chan/encrypt"
 	"honoka-chan/tools"
 	"honoka-chan/utils"
@@ -129,7 +129,7 @@ func Handshake(ctx *gin.Context) {
 	// fmt.Println(randKey)
 	// fmt.Println(deviceId)
 
-	err = database.LevelDb.Put([]byte(deviceId), []byte(randKey))
+	err = db.DB.Set([]byte(deviceId), []byte(randKey))
 	CheckErr(err)
 
 	token := strings.ToUpper(utils.RandomStr(33))
@@ -156,7 +156,7 @@ func Initialize(ctx *gin.Context) {
 	// fmt.Println(string(body64))
 
 	deviceId := ctx.Request.Header.Get("X-DEVICEID")
-	randKey, err := database.LevelDb.Get([]byte(deviceId))
+	randKey, err := db.DB.Get([]byte(deviceId))
 	CheckErr(err)
 	// decryptedBody, err := openssl.Des3ECBDecrypt(body64, randKey[0:24], openssl.PKCS7_PADDING)
 	// CheckErr(err)
@@ -205,7 +205,7 @@ func LoginAuto(ctx *gin.Context) {
 	// fmt.Println(string(body64))
 
 	deviceId := ctx.Request.Header.Get("X-DEVICEID")
-	randKey, err := database.LevelDb.Get([]byte(deviceId))
+	randKey, err := db.DB.Get([]byte(deviceId))
 	CheckErr(err)
 
 	decryptedBody, err := openssl.Des3ECBDecrypt(body64, randKey[0:24], openssl.PKCS7_PADDING)
@@ -279,7 +279,7 @@ func AccountLogin(ctx *gin.Context) {
 	// fmt.Println(string(body64))
 
 	deviceId := ctx.Request.Header.Get("X-DEVICEID")
-	randKey, err := database.LevelDb.Get([]byte(deviceId))
+	randKey, err := db.DB.Get([]byte(deviceId))
 	CheckErr(err)
 
 	decryptedBody, err := openssl.Des3ECBDecrypt(body64, randKey[0:24], openssl.PKCS7_PADDING)
@@ -432,7 +432,7 @@ func ReportRole(ctx *gin.Context) {
 	// fmt.Println(string(body64))
 
 	deviceId := ctx.Request.Header.Get("X-DEVICEID")
-	randKey, err := database.LevelDb.Get([]byte(deviceId))
+	randKey, err := db.DB.Get([]byte(deviceId))
 	CheckErr(err)
 
 	// decryptedBody, err := openssl.Des3ECBDecrypt(body64, randKey[0:24], openssl.PKCS7_PADDING)

main.go

@@ -1,13 +1,29 @@
 package main
 
 import (
+	"honoka-chan/db"
 	"honoka-chan/router"
 	_ "honoka-chan/tools"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/gin-gonic/gin"
 )
 
 func main() {
+	// 处理系统信号，确保程序退出时关闭数据库
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		<-signalChan
+		log.Println("Shutting down...")
+		db.DB.Close()
+		os.Exit(0)
+	}()
+
 	// Gin
 	gin.SetMode(gin.ReleaseMode)
 

middleware/authkey.go

@@ -4,7 +4,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"honoka-chan/encrypt"
 	"honoka-chan/utils"
 	"time"
@@ -37,7 +37,7 @@ func AuthKey(ctx *gin.Context) {
 		"server_token": serverToken,
 	})
 	CheckErr(err)
-	err = database.LevelDb.Put([]byte(authorizeToken), authJson)
+	err = db.DB.Set([]byte(authorizeToken), authJson)
 	CheckErr(err)
 
 	nonce := ctx.GetInt("nonce")

middleware/common.go

@@ -2,7 +2,7 @@ package middleware
 
 import (
 	"fmt"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -54,14 +54,14 @@ func Common(ctx *gin.Context) {
 		}
 		ctx.Set("userid", userId)
 
-		rToken, err := database.LevelDb.Get([]byte(userId))
+		rToken, err := db.DB.Get([]byte(userId))
 		CheckErr(err)
 		if token != string(rToken) {
 			ctx.String(http.StatusForbidden, ErrorMsg)
 			ctx.Abort()
 		}
 
-		if !database.MatchTokenUid(token, userId) {
+		if !db.MatchTokenUid(token, userId) {
 			ctx.String(http.StatusForbidden, ErrorMsg)
 			ctx.Abort()
 		}

middleware/login.go

@@ -2,7 +2,7 @@ package middleware
 
 import (
 	"encoding/base64"
-	"honoka-chan/database"
+	"honoka-chan/db"
 	"honoka-chan/encrypt"
 	"honoka-chan/utils"
 
@@ -11,7 +11,7 @@ import (
 )
 
 func Login(ctx *gin.Context) {
-	authData, err := database.LevelDb.Get([]byte(ctx.GetString("token")))
+	authData, err := db.DB.Get([]byte(ctx.GetString("token")))
 	CheckErr(err)
 
 	clientToken, err := base64.StdEncoding.DecodeString(gjson.Get(string(authData), "client_token").String())