modelRT/util/token.go

// 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, secretKey string) (string, error) {
	finalSecretKey := secretKey
	if finalSecretKey == "" {
		finalSecretKey = os.Getenv("TOKEN_SECRET_KEY")
	}

	if finalSecretKey == "" {
		return "", fmt.Errorf("TOKEN_SECRET_KEY environment variable not set and no key provided in parameters")
	}

	uniqueID := fmt.Sprintf("%d", time.Now().UnixNano())
	clientInfo := fmt.Sprintf("host=%s;service=%s;id=%s", host, serviceName, uniqueID)

	mac := hmac.New(sha256.New, []byte(finalSecretKey))
	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
}