leveldb存储token的简单实现

在之前的文章中leveldb的grpc接口没有实现ttl，这期加上。

一、改造leveldb的grpc接口支持ttl

这里简单改造，value值多存一个时间值，查的时候比较这个值即可

leveldb.proto

syntax = "proto3";option go_package = "../src/leveldb;leveldb";
package leveldb;service LevelDBService {rpc Put(PutRequest) returns (PutResponse) {}rpc Get(GetRequest) returns (GetResponse) {}rpc Has(HasRequest) returns (HasResponse) {}rpc Delete(DeleteRequest) returns (DeleteResponse) {}rpc GetTTL(GetTTLRequest) returns (GetTTLResponse) {}  // 新增方法
}message PutRequest {string key = 1;string value = 2;int64 ttl = 3;  // 单位：秒
}message PutResponse {}message GetRequest {string key = 1;
}message GetResponse {string value = 1;
}message HasRequest {string key = 1;
}message HasResponse {bool exists = 1;
}message DeleteRequest {string key = 1;
}message DeleteResponse {}message GetTTLRequest {string key = 1;
}message GetTTLResponse {int64 ttl = 1;  // 单位：秒
}

main.go

package mainimport ("context""encoding/binary""log""net""time"pb "gin/src/leveldb" // 替换为实际路径leveldb "github.com/syndtr/goleveldb/leveldb""google.golang.org/grpc""google.golang.org/grpc/reflection"
)// server 结构体嵌入 UnimplementedLevelDBServiceServer
type server struct {pb.UnimplementedLevelDBServiceServerdb *leveldb.DB
}func (s *server) Put(ctx context.Context, req *pb.PutRequest) (*pb.PutResponse, error) {key := []byte(req.GetKey())value := []byte(req.GetValue())// 处理 TTLttl := time.Duration(req.GetTtl()) * time.Secondif ttl > 0 {// 设置过期时间expirationTime := time.Now().Add(ttl).UnixNano()expirationData := make([]byte, 8)binary.BigEndian.PutUint64(expirationData, uint64(expirationTime))value = append(value, expirationData...)}err := s.db.Put(key, value, nil)if err != nil {return nil, err}return &pb.PutResponse{}, nil
}func (s *server) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetResponse, error) {key := []byte(req.GetKey())value, err := s.db.Get(key, nil)if err == leveldb.ErrNotFound {return &pb.GetResponse{Value: ""}, nil} else if err != nil {return nil, err}// 检查过期时间if len(value) > 8 {expirationTime := int64(binary.BigEndian.Uint64(value[len(value)-8:]))if time.Now().UnixNano() > expirationTime {return &pb.GetResponse{Value: ""}, nil}value = value[:len(value)-8]}return &pb.GetResponse{Value: string(value)}, nil
}func (s *server) Has(ctx context.Context, req *pb.HasRequest) (*pb.HasResponse, error) {key := []byte(req.GetKey())value, err := s.db.Get(key, nil)if err == leveldb.ErrNotFound {return &pb.HasResponse{Exists: false}, nil} else if err != nil {return nil, err}// 检查过期时间if len(value) > 8 {expirationTime := int64(binary.BigEndian.Uint64(value[len(value)-8:]))if time.Now().UnixNano() > expirationTime {return &pb.HasResponse{Exists: false}, nil}}return &pb.HasResponse{Exists: true}, nil
}func (s *server) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.DeleteResponse, error) {key := []byte(req.GetKey())err := s.db.Delete(key, nil)if err != nil {return nil, err}return &pb.DeleteResponse{}, nil
}func (s *server) GetTTL(ctx context.Context, req *pb.GetTTLRequest) (*pb.GetTTLResponse, error) {key := []byte(req.GetKey())value, err := s.db.Get(key, nil)if err == leveldb.ErrNotFound {return &pb.GetTTLResponse{Ttl: -1}, nil // 返回 -1 表示键不存在} else if err != nil {return nil, err}// 检查过期时间if len(value) > 8 {expirationTime := int64(binary.BigEndian.Uint64(value[len(value)-8:]))currentTime := time.Now().UnixNano()if currentTime > expirationTime {return &pb.GetTTLResponse{Ttl: 0}, nil // 返回 0 表示已过期}remainingTTL := (expirationTime - currentTime) / int64(time.Second)return &pb.GetTTLResponse{Ttl: remainingTTL}, nil}return &pb.GetTTLResponse{Ttl: -1}, nil // 返回 -1 表示没有设置 TTL
}func main() {// 打开或创建一个新的LevelDB数据库dbPath := "./data/leveldb/"db, err := leveldb.OpenFile(dbPath, nil)if err != nil {log.Fatalf("Failed to open database: %v", err)}defer db.Close()lis, err := net.Listen("tcp", ":3051")if err != nil {log.Fatalf("failed to listen: %v", err)}s := grpc.NewServer()pb.RegisterLevelDBServiceServer(s, &server{db: db})// 启用服务器反射reflection.Register(s)// 在启动时输出一条日志信息log.Println("gRPC server started on :3051")if err := s.Serve(lis); err != nil {log.Fatalf("failed to serve: %v", err)}
}

其他代码不变。

二、实现token认证工具类

leveldbTokenOperator.go

package middlewareimport "time"import ("context""encoding/json""gin-epg/src/leveldb" // 假设 gRPC 服务的 proto 文件已经编译并生成了 Go 代码"github.com/google/uuid""google.golang.org/grpc""google.golang.org/grpc/codes""google.golang.org/grpc/status"
)const (TokenPrefixKey         = "token:"UserPrefixKey          = "user:"SingleUserMaxTokenSize = 10defaultExpirationTime  = 1209600 // 2 weeks in seconds
)type LeveldbSimpleTokenOperator struct {expirationTimeInSecond int64client                 leveldb.LevelDBServiceClient
}func NewLeveldbSimpleTokenOperator(addr string, expirationTimeInSecond ...int64) (*LeveldbSimpleTokenOperator, error) {conn, err := grpc.Dial(addr, grpc.WithInsecure())if err != nil {return nil, err}expiration := defaultExpirationTimeif len(expirationTimeInSecond) > 0 {expiration = int(expirationTimeInSecond[0])}client := leveldb.NewLevelDBServiceClient(conn)return &LeveldbSimpleTokenOperator{expirationTimeInSecond: int64(expiration),client:                 client,}, nil
}// getUserMapFromToken 从 token 中获取 UserMap
func (r *LeveldbSimpleTokenOperator) getUserMapFromToken(token string) (map[string]interface{}, error) {key := TokenPrefixKey + tokenresp, err := r.client.Get(context.Background(), &leveldb.GetRequest{Key: key})if err != nil {if status.Code(err) == codes.NotFound {return nil, nil}return nil, err}var userMap map[string]interface{}err = json.Unmarshal([]byte(resp.Value), &userMap)if err != nil {return nil, err}return userMap, nil
}// getExpirationDateFromToken 从 token 中获取过期日
func (r *LeveldbSimpleTokenOperator) getExpirationDateFromToken(token string) (time.Time, error) {resp, err := r.client.GetTTL(context.Background(), &leveldb.GetTTLRequest{Key: token})if err != nil {return time.Time{}, err}return time.Now().Add(time.Duration(resp.Ttl) * time.Second), nil
}// isTokenExpired 判断 token 是否过期
func (r *LeveldbSimpleTokenOperator) isTokenExpired(token string) (bool, error) {key := TokenPrefixKey + tokenresp, err := r.client.Has(context.Background(), &leveldb.HasRequest{Key: key})if err != nil {return false, err}return !resp.Exists, nil
}// generateToken 生成 token
func (r *LeveldbSimpleTokenOperator) generateToken(userMap map[string]interface{}) (string, error) {token := uuid.New().String()key := TokenPrefixKey + tokenvalue, err := json.Marshal(userMap)if err != nil {return "", err}_, err = r.client.Put(context.Background(), &leveldb.PutRequest{Key: key, Value: string(value), Ttl: r.expirationTimeInSecond})if err != nil {return "", err}return token, nil
}// validateToken 判断 token 是否有效
func (r *LeveldbSimpleTokenOperator) validateToken(token string) (bool, error) {isExpired, err := r.isTokenExpired(token)if err != nil {return false, err}return !isExpired, nil
}

gin中间件封装

middlewareUtil.go

package middlewareimport ("fmt""gin-epg/src/common/util""github.com/gin-gonic/gin""net/http""strings"
)// TokenMiddleware 是一个用于校验 token 的中间件
func TokenMiddleware() gin.HandlerFunc {return func(c *gin.Context) {var token string// 尝试从 Authorization 请求头获取 tokenauthHeader := c.GetHeader("Authorization")if authHeader != "" {// 按空格分割parts := strings.SplitN(authHeader, " ", 2)if len(parts) == 2 && parts[0] == "Bearer" {token = parts[1]} else {c.JSON(http.StatusBadRequest, gin.H{"code": 2004,"msg":  "请求头中auth格式有误",})c.Abort()return}} else {// 如果 Authorization 头不存在，尝试从 token 请求头获取 tokentoken = c.GetHeader("token")if token == "" {c.JSON(http.StatusUnauthorized, gin.H{"error": "未提供 token"})c.Abort()return}}// 使用解析JWT的函数来解析 tokenclaims, err := ParseToken(token)if err != nil {c.JSON(http.StatusUnauthorized, gin.H{"code": 2005,"msg":  "无效的Token",})c.Abort()return}// 将用户信息设置到请求上下文中c.Set("id", claims["id"])c.Set("username", claims["username"])c.Set("role", claims["role"])c.Set("avatarUrl", claims["avatarUrl"])c.Next() // 后续的处理函数可以通过c.Get("username")等来获取当前请求的用户信息}
}// ParseToken 解析token并返回用户信息
func ParseToken(s string) (map[string]interface{}, error) {// 拼接 fileDownloadUrl 和 filePathaddr := "localhost:3051"configAddr, err := util.GetConfigValue("leveldbRpcUrl")if err == nil {addr = configAddr.(string) // 进行类型断言}operator, err := NewLeveldbSimpleTokenOperator(addr)if err != nil {return nil, err}claims, err := operator.getUserMapFromToken(s)if err != nil {return nil, err}return claims, nil
}// GenerateToken 生成Token
func GenerateToken(userMap map[string]interface{}) (string, error) {// 获取配置中的地址addr := "localhost:3051"configAddr, err := util.GetConfigValue("leveldbRpcUrl")if err == nil {addr = configAddr.(string) // 进行类型断言}// 创建LevelDB Token操作器operator, err := NewLeveldbSimpleTokenOperator(addr)if err != nil {return "", fmt.Errorf("failed to create LevelDB token operator: %w", err)}// 生成Tokenclaims, err := operator.generateToken(userMap)if err != nil {return "", fmt.Errorf("failed to generate token: %w", err)}return claims, nil
}

登录时生成token调用示例

package controllerimport ("crypto/md5""fmt""gin-epg/src/entity""gin-epg/src/middleware""gin-epg/src/service""github.com/gin-gonic/gin""net/http"
)// RestResponse 响应结构体
type RestResponse struct {Status  int         `json:"status"`Result  interface{} `json:"result,omitempty"`Message string      `json:"message,omitempty"`
}// LoginData 登录数据结构体
type LoginData struct {Username string `json:"username"`Password string `json:"password"`
}// UserInfo 用户信息结构体
type UserInfo struct {Token string       `json:"token"`User  *entity.User `json:"user"`
}// doLogin 登录处理函数
func DoLogin(c *gin.Context) {var loginData LoginDataif err := c.ShouldBindJSON(&loginData); err != nil {c.JSON(http.StatusBadRequest, RestResponse{Status: 400, Message: "请求参数错误"})return}username := loginData.Usernamepassword := loginData.PasswordloginUser, err := service.FindUserByName(username)if err != nil || loginUser == nil {c.JSON(http.StatusUnauthorized, RestResponse{Status: 401, Message: "用户不存在"})return}dbPassword := loginUser.PasswordencryptedText := MD5(password)if dbPassword != encryptedText {c.JSON(http.StatusUnauthorized, RestResponse{Status: 401, Message: "登录失败"})return}userInfoClaims := map[string]interface{}{"id":        loginUser.ID,"username":  loginUser.Username,"role":      "user","avatarUrl": loginUser.AvatarURL,}token, err := middleware.GenerateToken(userInfoClaims)if err != nil {c.JSON(http.StatusInternalServerError, RestResponse{Status: 500, Message: "生成Token失败"})return}userInfo := UserInfo{Token: token,User:  loginUser,}c.JSON(http.StatusOK, RestResponse{Status: 200, Result: userInfo})
}func MD5(str string) string {data := []byte(str) //切片has := md5.Sum(data)md5str := fmt.Sprintf("%x", has) //将[]byte转成16进制return md5str
}

调用接口校验token示例

单个接口使用

epgChannelGroup.GET("/deleteChannelByName", middleware.TokenMiddleware(), controller.DeleteEpgChannelByName)

group使用

epgChannelGroup.Use(middleware.TokenMiddleware())