From 9dcfc9022655e8cfe2b58fdd82f70031aa25f3b1 Mon Sep 17 00:00:00 2001 From: Yuan Si Date: Mon, 11 Nov 2024 20:11:05 +0800 Subject: [PATCH] Update leveldb Signed-off-by: Yuan Si --- database/leveldb.go | 68 ------------------------------------------- database/utils.go | 18 ------------ db/db.go | 21 +++++++++++++ db/leveldb.go | 58 ++++++++++++++++++++++++++++++++++++ handler/live.go | 6 ++-- handler/login.go | 4 +-- handler/private.go | 12 ++++---- main.go | 16 ++++++++++ middleware/authkey.go | 4 +-- middleware/common.go | 6 ++-- middleware/login.go | 4 +-- 11 files changed, 113 insertions(+), 104 deletions(-) delete mode 100644 database/leveldb.go delete mode 100644 database/utils.go create mode 100644 db/db.go create mode 100644 db/leveldb.go diff --git a/database/leveldb.go b/database/leveldb.go deleted file mode 100644 index 5549992..0000000 --- a/database/leveldb.go +++ /dev/null @@ -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 -} diff --git a/database/utils.go b/database/utils.go deleted file mode 100644 index 4426077..0000000 --- a/database/utils.go +++ /dev/null @@ -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 -} diff --git a/db/db.go b/db/db.go new file mode 100644 index 0000000..63ba8d8 --- /dev/null +++ b/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 +} diff --git a/db/leveldb.go b/db/leveldb.go new file mode 100644 index 0000000..b2d8e51 --- /dev/null +++ b/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) +} diff --git a/handler/live.go b/handler/live.go index 8ab6d93..89d5030 100644 --- a/handler/live.go +++ b/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"). diff --git a/handler/login.go b/handler/login.go index cdbf130..0265025 100644 --- a/handler/login.go +++ b/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{ diff --git a/handler/private.go b/handler/private.go index 6851b18..7536aec 100644 --- a/handler/private.go +++ b/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) diff --git a/main.go b/main.go index fd006f1..f27cb97 100644 --- a/main.go +++ b/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) diff --git a/middleware/authkey.go b/middleware/authkey.go index 3b5267d..f7e92e4 100644 --- a/middleware/authkey.go +++ b/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") diff --git a/middleware/common.go b/middleware/common.go index d964996..0550b69 100644 --- a/middleware/common.go +++ b/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() } diff --git a/middleware/login.go b/middleware/login.go index 586c0d7..9a5a265 100644 --- a/middleware/login.go +++ b/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())