101 lines
2.9 KiB
Go
101 lines
2.9 KiB
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) (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.")
|
||
}
|
||
}
|