add func of generate service token

config/config.go

@@ -7,14 +7,20 @@ import (
 	"github.com/spf13/viper"
 )
 
-// BaseConfig define config stuct of base params config
+// BaseConfig define config struct of base params config
 type BaseConfig struct {
 	GridID    int64 `mapstructure:"grid_id"`
 	ZoneID    int64 `mapstructure:"zone_id"`
 	StationID int64 `mapstructure:"station_id"`
 }
 
-// KafkaConfig define config stuct of kafka config
+// ServiceConfig define config struct of service config
+type ServiceConfig struct {
+	ServiceName string `mapstructure:"service_name"`
+	SecretKey   string `mapstructure:"secret_key"`
+}
+
+// KafkaConfig define config struct of kafka config
 type KafkaConfig struct {
 	Servers                 string `mapstructure:"Servers"`
 	GroupID                 string `mapstructure:"group_id"`
@@ -24,7 +30,7 @@ type KafkaConfig struct {
 	ReadMessageTimeDuration string `mapstructure:"read_message_time_duration"`
 }
 
-// PostgresConfig define config stuct of postgres config
+// PostgresConfig define config struct of postgres config
 type PostgresConfig struct {
 	Port     int    `mapstructure:"port"`
 	Host     string `mapstructure:"host"`
@@ -33,7 +39,7 @@ type PostgresConfig struct {
 	Password string `mapstructure:"password"`
 }
 
-// LoggerConfig define config stuct of zap logger config
+// LoggerConfig define config struct of zap logger config
 type LoggerConfig struct {
 	Mode       string `mapstructure:"mode"`
 	Level      string `mapstructure:"level"`
@@ -44,7 +50,7 @@ type LoggerConfig struct {
 	Compress   bool   `mapstructure:"compress"`
 }
 
-// RedisConfig define config stuct of redis config
+// RedisConfig define config struct of redis config
 type RedisConfig struct {
 	Addr     string `mapstructure:"addr"`
 	Password string `mapstructure:"password"`
@@ -53,13 +59,13 @@ type RedisConfig struct {
 	Timeout  int    `mapstructure:"timeout"`
 }
 
-// AntsConfig define config stuct of ants pool config
+// AntsConfig define config struct of ants pool config
 type AntsConfig struct {
 	ParseConcurrentQuantity      int `mapstructure:"parse_concurrent_quantity"`       // parse comtrade file concurrent quantity
 	RTDReceiveConcurrentQuantity int `mapstructure:"rtd_receive_concurrent_quantity"` // polling real time data concurrent quantity
 }
 
-// DataRTConfig define config stuct of data runtime server api config
+// DataRTConfig define config struct of data runtime server api config
 type DataRTConfig struct {
 	Host       string `mapstructure:"host"`
 	Port       int64  `mapstructure:"port"`
@@ -67,9 +73,10 @@ type DataRTConfig struct {
 	Method     string `mapstructure:"polling_api_method"`
 }
 
-// ModelRTConfig define config stuct of model runtime server
+// ModelRTConfig define config struct of model runtime server
 type ModelRTConfig struct {
 	BaseConfig         `mapstructure:"base"`
+	ServiceConfig      `mapstructure:"service"`
 	PostgresConfig     `mapstructure:"postgres"`
 	KafkaConfig        `mapstructure:"kafka"`
 	LoggerConfig       `mapstructure:"logger"`

config/config.yaml

@@ -57,6 +57,11 @@ base:
   zone_id: 1
   station_id: 1
 
+# modelRT service config
+service:
+  service_name: "modelRT"
+  secret_key: "modelRT"
+
 # dataRT api config
 dataRT:
   host: "http://127.0.0.1"

handler/measurement_load.go

@@ -1,3 +1,4 @@
+// Package handler provides HTTP handlers for various endpoints.
 package handler
 
 import (
@@ -14,6 +15,23 @@ import (
 // MeasurementGetHandler retrieves the value of measurement data
 func MeasurementGetHandler(c *gin.Context) {
 	var request network.MeasurementGetRequest
+	var clientToken string
+
+	// TODO 优化client_token获取
+	token, ok := c.Get("client_token")
+	if ok {
+		// 将 token 断言为 ClientToken 类型
+		clientToken, isClientToken := token.(string)
+		if !isClientToken {
+			// TODO 优化相关报错输出
+			c.JSON(http.StatusOK, gin.H{
+				"message": "Request processed successfully",
+				"token":   clientToken,
+			})
+			return
+		}
+		clientToken = token.(string)
+	}
 
 	if err := c.ShouldBindJSON(&request); err != nil {
 		logger.Error(c, "Failed to unmarshal measurement get request", "error", err)
@@ -23,9 +41,8 @@ func MeasurementGetHandler(c *gin.Context) {
 		})
 		return
 	}
-
 	// token 当前客户端的唯一标识（token）,用于区分不同的客户端。
-	zset := diagram.NewRedisZSet(c, request.MeasurementToken, "token", 10, false)
+	zset := diagram.NewRedisZSet(c, request.MeasurementToken, clientToken, 0, false)
 	points, err := zset.ZRANGE(request.MeasurementToken, 0, -1)
 	if err != nil {
 		logger.Error(c, "failed to get measurement data from redis", "measurement_token", request.MeasurementToken, "error", err)

main.go

@@ -21,6 +21,7 @@ import (
 	"modelRT/middleware"
 	"modelRT/pool"
 	"modelRT/router"
+	"modelRT/util"
 
 	realtimedata "modelRT/real-time-data"
 
@@ -58,6 +59,19 @@ func main() {
 	logger.InitLoggerInstance(modelRTConfig.LoggerConfig)
 
 	modelRTConfig = config.ReadAndInitConfig(*modelRTConfigDir, *modelRTConfigName, *modelRTConfigType)
+
+	hostName, err := os.Hostname()
+	if err != nil {
+		logger.Error(ctx, "get host name failed", "error", err)
+		panic(err)
+	}
+
+	serviceToken, err := util.GenerateClientToken(hostName, modelRTConfig.ServiceConfig.ServiceName)
+	if err != nil {
+		logger.Error(ctx, "generate client token failed", "error", err)
+		panic(err)
+	}
+
 	// init postgresDBClient
 	postgresDBClient = database.InitPostgresDBInstance(ctx, modelRTConfig.PostgresDBURI)
 
@@ -123,7 +137,7 @@ func main() {
 	// go realtimedata.RealTimeDataComputer(ctx, nil, []string{}, "")
 
 	engine := gin.New()
-	router.RegisterRoutes(engine)
+	router.RegisterRoutes(engine, serviceToken)
 	server := http.Server{
 		Addr:    ":8080",
 		Handler: engine,

middleware/token.go

@@ -0,0 +1,11 @@
+package middleware
+
+import "github.com/gin-gonic/gin"
+
+// SetTokenMiddleware define a middleware for set token in context
+func SetTokenMiddleware(clientToken string) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		c.Set("client_token", clientToken)
+		c.Next()
+	}
+}

network/measurement_request.go

@@ -3,5 +3,6 @@ package network
 
 // MeasurementGetRequest defines the request payload for getting an measurement
 type MeasurementGetRequest struct {
+	MeasurementID    int64  `json:"measurement_id"`
 	MeasurementToken string `json:"token"`
 }

router/measurement.go

@@ -3,12 +3,14 @@ package router
 
 import (
 	"modelRT/handler"
+	"modelRT/middleware"
 
 	"github.com/gin-gonic/gin"
 )
 
 // registerMeasurementRoutes define func of register measurement routes
-func registerMeasurementRoutes(rg *gin.RouterGroup) {
+func registerMeasurementRoutes(rg *gin.RouterGroup, clientToken string) {
 	g := rg.Group("/measurement/")
+	rg.Use(middleware.SetTokenMiddleware(clientToken))
 	g.GET("load", handler.MeasurementGetHandler)
 }

router/router.go

@@ -16,11 +16,11 @@ func init() {
 }
 
 // RegisterRoutes define func of register all routes
-func RegisterRoutes(engine *gin.Engine) {
+func RegisterRoutes(engine *gin.Engine, clientToken string) {
 	// use global middlewares
 	engine.Use(middleware.StartTrace(), limiter.Middleware)
 	routeGroup := engine.Group("")
 	registerDiagramRoutes(routeGroup)
 	registerAttrRoutes(routeGroup)
-	registerMeasurementRoutes(routeGroup)
+	registerMeasurementRoutes(routeGroup, clientToken)
 }

util/redis_init.go

@@ -1,3 +1,4 @@
+// Package util provide some utility functions
 package util
 
 import (

util/redis_options.go

@@ -1,3 +1,4 @@
+// Package util provide some utility functions
 package util
 
 import (

util/token.go

@@ -0,0 +1,100 @@
+// Package util provide some utility functions
+package util
+
+import (
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"os"
+	"strings"
+	"time"
+)
+
+// GenerateClientToken define func of generate a secure token for client identification
+func GenerateClientToken(host string, serviceName string) (string, error) {
+	uniqueID := fmt.Sprintf("%d", time.Now().UnixNano())
+	clientInfo := fmt.Sprintf("host=%s;service=%s;id=%s", host, serviceName, uniqueID)
+	secretKey := os.Getenv("TOKEN_SECRET_KEY")
+	if secretKey == "" {
+		return "", fmt.Errorf("TOKEN_SECRET_KEY environment variable not set")
+	}
+
+	mac := hmac.New(sha256.New, []byte(secretKey))
+	mac.Write([]byte(clientInfo))
+	signature := mac.Sum(nil)
+
+	token := fmt.Sprintf("%s.%s",
+		base64.URLEncoding.EncodeToString([]byte(clientInfo)),
+		base64.URLEncoding.EncodeToString(signature))
+
+	return token, nil
+}
+
+// verifyClientToken define func of verify the client token
+func verifyClientToken(token string, secretKey string) (bool, string, error) {
+	parts := strings.Split(token, ".")
+	if len(parts) != 2 {
+		return false, "", fmt.Errorf("invalid token format")
+	}
+
+	encodedClientInfo := parts[0]
+	encodedSignature := parts[1]
+
+	clientInfoBytes, err := base64.URLEncoding.DecodeString(encodedClientInfo)
+	if err != nil {
+		return false, "", fmt.Errorf("failed to decode client info: %w", err)
+	}
+	signatureBytes, err := base64.URLEncoding.DecodeString(encodedSignature)
+	if err != nil {
+		return false, "", fmt.Errorf("failed to decode signature: %w", err)
+	}
+
+	mac := hmac.New(sha256.New, []byte(secretKey))
+	mac.Write(clientInfoBytes)
+	expectedSignature := mac.Sum(nil)
+
+	if !hmac.Equal(signatureBytes, expectedSignature) {
+		return false, "", nil // 签名不匹配，token 无效
+	}
+
+	clientInfo := string(clientInfoBytes)
+	return true, clientInfo, nil
+}
+
+// --- 示例使用 ---
+
+func main() {
+	// 模拟获取客户端信息
+	// 实际应用中，host 可以通过获取本地主机名，serviceName 可以是当前服务的标识。
+	host, err := os.Hostname()
+	if err != nil {
+		fmt.Println("Error getting hostname:", err)
+		host = "unknown-host"
+	}
+	serviceName := "my-web-service" // 替换为你的服务名
+
+	// 设置一个示例密钥（在实际应用中，请从环境变量加载）
+	// export TOKEN_SECRET_KEY="your-super-secret-key-here"
+	os.Setenv("TOKEN_SECRET_KEY", "my-super-secret-and-long-key-for-hmac")
+
+	token, err := GenerateClientToken(host, serviceName)
+	if err != nil {
+		fmt.Println("Error generating token:", err)
+		return
+	}
+
+	fmt.Println("Generated Token:", token)
+
+	// --- 验证 Token (示例) ---
+	// 在接收方，你需要实现一个验证函数
+	isValid, clientInfo, err := verifyClientToken(token, os.Getenv("TOKEN_SECRET_KEY"))
+	if err != nil {
+		fmt.Println("Error verifying token:", err)
+	} else if isValid {
+		fmt.Println("Token is valid!")
+		fmt.Println("Client Info:", clientInfo) // 包含 host, service, id
+	} else {
+		fmt.Println("Token is invalid.")
+	}
+}

util/trace.go

@@ -1,3 +1,4 @@
+// Package util provide some utility functions
 package util
 
 import (