Files
waf-platform/EdgeNode/internal/nodes/http_request_token.go
2026-02-04 20:27:13 +08:00

121 lines
2.8 KiB
Go

// 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()
}