1.4.5.2

EdgeNode/internal/nodes/http_request_token.go

@@ -0,0 +1,120 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package nodes
+
+import (
+	"encoding/json"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/ttlcache"
+)
+
+// TokenCache Token 缓存
+type TokenCache struct {
+	cache *ttlcache.Cache[string] // IP => Token
+}
+
+var sharedTokenCache *TokenCache
+var sharedTokenCacheOnce sync.Once
+
+// SharedTokenCache 获取共享 Token 缓存实例
+func SharedTokenCache() *TokenCache {
+	sharedTokenCacheOnce.Do(func() {
+		sharedTokenCache = NewTokenCache()
+	})
+	return sharedTokenCache
+}
+
+// NewTokenCache 创建新 Token 缓存
+func NewTokenCache() *TokenCache {
+	cache := ttlcache.NewCache[string](
+		ttlcache.NewMaxItemsOption(10000),
+		ttlcache.NewGCConfigOption().
+			WithBaseInterval(5*time.Minute).
+			WithMinInterval(2*time.Minute).
+			WithMaxInterval(10*time.Minute).
+			WithAdaptive(true),
+	)
+
+	return &TokenCache{
+		cache: cache,
+	}
+}
+
+// Set 设置 Token
+func (this *TokenCache) Set(ip string, token string) {
+	expiresAt := fasttime.Now().Unix() + 300 // 5 分钟过期
+	this.cache.Write(ip, token, expiresAt)
+}
+
+// Get 获取 Token
+func (this *TokenCache) Get(ip string) (string, bool) {
+	item := this.cache.Read(ip)
+	if item == nil {
+		return "", false
+	}
+
+	if item.ExpiresAt() < fasttime.Now().Unix() {
+		return "", false
+	}
+
+	return item.Value, true
+}
+
+// ValidateToken 验证 Token
+func (this *HTTPRequest) validateEncryptionToken(token string) bool {
+	if this.web.Encryption == nil || !this.web.Encryption.IsOn || !this.web.Encryption.IsEnabled() {
+		return true // 未启用加密，直接通过
+	}
+
+	if len(token) == 0 {
+		return false
+	}
+
+	remoteIP := this.requestRemoteAddr(true)
+	cache := SharedTokenCache()
+	storedToken, ok := cache.Get(remoteIP)
+	if !ok {
+		return false
+	}
+
+	return storedToken == token
+}
+
+// handleTokenHandshake 处理 Token 握手请求
+func (this *HTTPRequest) handleTokenHandshake() {
+	if this.RawReq.Method != http.MethodPost {
+		this.writeCode(http.StatusMethodNotAllowed, "Method Not Allowed", "方法不允许")
+		return
+	}
+
+	// 解析请求体
+	var requestBody struct {
+		Token string `json:"token"`
+	}
+
+	err := json.NewDecoder(this.RawReq.Body).Decode(&requestBody)
+	if err != nil {
+		this.writeCode(http.StatusBadRequest, "Invalid Request", "请求格式错误")
+		return
+	}
+
+	if len(requestBody.Token) == 0 {
+		this.writeCode(http.StatusBadRequest, "Token Required", "Token 不能为空")
+		return
+	}
+
+	// 保存 Token
+	remoteIP := this.requestRemoteAddr(true)
+	cache := SharedTokenCache()
+	cache.Set(remoteIP, requestBody.Token)
+
+	// 返回成功响应
+	this.writer.Header().Set("Content-Type", "application/json")
+	this.writer.WriteHeader(http.StatusOK)
+	_, _ = this.writer.WriteString(`{"success":true}`)
+	this.writer.SetOk()
+}