diff --git a/internal/apps/openflare/agent/auth_cache.go b/internal/apps/openflare/agent/auth_cache.go index 42b8d487..a52061ae 100644 --- a/internal/apps/openflare/agent/auth_cache.go +++ b/internal/apps/openflare/agent/auth_cache.go @@ -19,6 +19,8 @@ import ( const ( agentTokenPositiveCacheTTL = 2 * time.Minute agentTokenNegativeCacheTTL = 10 * time.Minute + // ponytail: 上限仅防未授权口伪造 token 撑爆内存;打满后放弃缓存(回退 DB 查询),行为不变 + maxAgentTokenNegativeCacheEntries = 10_000 ) type cachedAgentNode struct { @@ -118,9 +120,24 @@ func (c *accessTokenAuthCache) storeMissing(token string, expiresAt time.Time) { c.mu.Lock() defer c.mu.Unlock() delete(c.positive, token) + if len(c.negative) >= maxAgentTokenNegativeCacheEntries { + c.evictExpiredMissingLocked(c.now()) + if len(c.negative) >= maxAgentTokenNegativeCacheEntries { + return // 缓存满:放弃缓存该 token,认证仍走 DB,仅防内存无限增长 + } + } c.negative[token] = expiresAt } +// evictExpiredMissingLocked 清理已过期的 negative 条目,须持写锁调用。 +func (c *accessTokenAuthCache) evictExpiredMissingLocked(now time.Time) { + for token, expiresAt := range c.negative { + if now.After(expiresAt) { + delete(c.negative, token) + } + } +} + func (c *accessTokenAuthCache) reset() { c.mu.Lock() defer c.mu.Unlock() diff --git a/internal/apps/openflare/flared/middleware.go b/internal/apps/openflare/flared/middleware.go index 1e800efe..3d5bb9f6 100644 --- a/internal/apps/openflare/flared/middleware.go +++ b/internal/apps/openflare/flared/middleware.go @@ -4,16 +4,13 @@ package flared import ( - "context" - "errors" "strings" - "github.com/Rain-kl/Wavelet/internal/repository" + "github.com/Rain-kl/Wavelet/internal/apps/openflare/agent" - "github.com/Rain-kl/Wavelet/internal/model" "github.com/Rain-kl/Wavelet/internal/shared/response" "github.com/gin-gonic/gin" - "gorm.io/gorm" + ) const ctxFlaredNodeKey = "flared_node" @@ -22,7 +19,7 @@ const ctxFlaredNodeKey = "flared_node" func TunnelAuth() gin.HandlerFunc { return func(c *gin.Context) { token := strings.TrimSpace(c.GetHeader("X-Tunnel-Token")) - node, err := authenticateAccessToken(c.Request.Context(), token) + node, err := agent.AuthenticateAccessToken(c.Request.Context(), token) if err != nil { response.AbortUnauthorized(c, errTunnelTokenInvalid) return @@ -35,17 +32,3 @@ func TunnelAuth() gin.HandlerFunc { c.Next() } } - -func authenticateAccessToken(ctx context.Context, token string) (*model.OpenFlareNode, error) { - if token == "" { - return nil, errors.New("missing tunnel token") - } - node, err := repository.GetOpenFlareNodeByAccessToken(ctx, token) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, errors.New("invalid tunnel token") - } - return nil, err - } - return node, nil -} diff --git a/internal/apps/openflare/relay/middleware.go b/internal/apps/openflare/relay/middleware.go index b0e3d10e..18f14d56 100644 --- a/internal/apps/openflare/relay/middleware.go +++ b/internal/apps/openflare/relay/middleware.go @@ -4,16 +4,13 @@ package relay import ( - "context" - "errors" "strings" - "github.com/Rain-kl/Wavelet/internal/repository" + "github.com/Rain-kl/Wavelet/internal/apps/openflare/agent" - "github.com/Rain-kl/Wavelet/internal/model" "github.com/Rain-kl/Wavelet/internal/shared/response" "github.com/gin-gonic/gin" - "gorm.io/gorm" + ) const ctxRelayNodeKey = "relay_node" @@ -22,7 +19,7 @@ const ctxRelayNodeKey = "relay_node" func Auth() gin.HandlerFunc { return func(c *gin.Context) { token := strings.TrimSpace(c.GetHeader("X-Agent-Token")) - node, err := authenticateAccessToken(c.Request.Context(), token) + node, err := agent.AuthenticateAccessToken(c.Request.Context(), token) if err != nil { response.AbortUnauthorized(c, errAgentTokenInvalid) return @@ -35,17 +32,3 @@ func Auth() gin.HandlerFunc { c.Next() } } - -func authenticateAccessToken(ctx context.Context, token string) (*model.OpenFlareNode, error) { - if token == "" { - return nil, errors.New("missing agent token") - } - node, err := repository.GetOpenFlareNodeByAccessToken(ctx, token) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, errors.New("invalid agent token") - } - return nil, err - } - return node, nil -}