From 2d3dde99c2c7dd3defed7e5358a4944a6c90ee96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=AD=E5=B7=9E=E6=98=8E=E5=A9=B3=E7=A7=91=E6=8A=80?= Date: Tue, 7 Jul 2026 00:13:56 +0800 Subject: [PATCH] fix: harden jwt locks sse and storage safety --- cache/lock.go | 39 ++++++++-- cache/lock_concurrency_test.go | 79 ++++++++++++++++++++ cache/lock_test.go | 6 +- compress/compress.go | 2 +- jwt/jwt.go | 119 +++++++++++++++++++----------- jwt/jwt_test.go | 129 ++++++++++++++++++++++++++++++++- sse/sse.go | 23 +++++- sse/sse_test.go | 47 +++++++++++- storage/storage.go | 33 +++++++-- 9 files changed, 414 insertions(+), 63 deletions(-) diff --git a/cache/lock.go b/cache/lock.go index 3c2518d..03e4cba 100644 --- a/cache/lock.go +++ b/cache/lock.go @@ -15,6 +15,12 @@ var ( ErrLockNotHeld = errors.New("锁未被当前客户端持有") ErrLockExpired = errors.New("锁已过期") ErrRedisNotReady = errors.New("Redis 未初始化") + // ErrLockNotAcquired 表示锁被其它客户端持有，业务函数未执行。 + ErrLockNotAcquired = errors.New("未获取到锁") + // ErrInvalidLockTTL 表示锁 TTL 小于 Redis PX 支持的 1ms 粒度或非正。 + ErrInvalidLockTTL = errors.New("锁 TTL 必须大于等于 1ms") + // ErrInvalidLockRetryInterval 表示重试或续期间隔非法。 + ErrInvalidLockRetryInterval = errors.New("锁重试/续期间隔必须大于 0") // ErrLockUnexpectedResult Lua 脚本返回了非预期的结果类型（C1b：裸类型断言防护）。 ErrLockUnexpectedResult = errors.New("锁脚本返回非预期结果") ) @@ -29,6 +35,13 @@ func toInt64(v any) (int64, error) { return n, nil } +func ttlMillis(ttl time.Duration) (int64, error) { + if ttl < time.Millisecond { + return 0, ErrInvalidLockTTL + } + return int64(ttl / time.Millisecond), nil +} + // LockToken 锁令牌（用于安全释放锁）。 // // 安全说明（C1d 设计局限）：Token 是随机 UUID（非单调递增的 fencing token）。 @@ -88,7 +101,10 @@ func NewLock(ctx context.Context, key string, ttl time.Duration) (*LockToken, er } token := utils.UUID() - ttlMs := int64(ttl / time.Millisecond) + ttlMs, err := ttlMillis(ttl) + if err != nil { + return nil, err + } result, err := rdb.Eval(ctx, lockScript, []string{key}, token, ttlMs).Result() if err != nil { @@ -148,7 +164,7 @@ func Unlock(ctx context.Context, token *LockToken) error { func UnlockByKey(ctx context.Context, key string) error { rdb := database.GetRedis() if rdb == nil { - return nil + return ErrRedisNotReady } return rdb.Del(ctx, key).Err() } @@ -165,7 +181,10 @@ func ExtendLock(ctx context.Context, token *LockToken, ttl time.Duration) error return ErrLockNotHeld } - ttlMs := int64(ttl / time.Millisecond) + ttlMs, err := ttlMillis(ttl) + if err != nil { + return err + } result, err := rdb.Eval(ctx, extendScript, []string{token.Key}, token.Token, ttlMs).Result() if err != nil { @@ -185,6 +204,9 @@ func ExtendLock(ctx context.Context, token *LockToken, ttl time.Duration) error // TryLock 尝试获取锁，失败时等待重试。重试等待响应 ctx 取消（C1c 修复）。 func TryLock(ctx context.Context, key string, ttl time.Duration, retryInterval time.Duration, maxRetry int) (*LockToken, error) { + if retryInterval <= 0 { + return nil, ErrInvalidLockRetryInterval + } for i := 0; i < maxRetry; i++ { token, err := NewLock(ctx, key, ttl) if err != nil { @@ -200,7 +222,7 @@ func TryLock(ctx context.Context, key string, ttl time.Duration, retryInterval t case <-time.After(retryInterval): } } - return nil, nil + return nil, ErrLockNotAcquired } // WithLock 使用分布式锁执行函数（自动管理锁）。 @@ -215,7 +237,7 @@ func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() erro return err } if token == nil { - return nil // 未获取到锁，跳过执行 + return ErrLockNotAcquired } defer func() { unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) @@ -235,12 +257,15 @@ func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() erro // （ctx 取消或 ExtendLock 失败时 done 已 closed，父再 send 即 panic，Unlock 不执行、锁泄漏到 TTL）。 // Unlock 用 context.Background() 派生超时，避免原 ctx 已取消致 Unlock 失败再泄漏。 func WithLockAutoExtend(ctx context.Context, key string, initialTTL time.Duration, extendInterval time.Duration, fn func() error) error { + if extendInterval <= 0 { + return ErrInvalidLockRetryInterval + } token, err := NewLock(ctx, key, initialTTL) if err != nil { return err } if token == nil { - return nil // 未获取到锁，跳过执行 + return ErrLockNotAcquired } // 父关停信号（仅父 close）与子 ack 信号（仅子 close）。 @@ -392,4 +417,4 @@ func SetRaw(ctx context.Context, key string, value string, ttl time.Duration) er return nil } return rdb.Set(ctx, key, value, ttl).Err() -} \ No newline at end of file +} diff --git a/cache/lock_concurrency_test.go b/cache/lock_concurrency_test.go index 466ccde..cc55590 100644 --- a/cache/lock_concurrency_test.go +++ b/cache/lock_concurrency_test.go @@ -299,3 +299,82 @@ func TestTryLockRespectsCtxCancel(t *testing.T) { t.Errorf("TryLock took %v, should return shortly after ctx cancel (C1c: time.Sleep ignored ctx)", elapsed) } } + +func TestLockRejectsSubMillisecondTTL(t *testing.T) { + setupMiniRedis(t) + ctx := context.Background() + + if token, err := cache.NewLock(ctx, "m10:ttl", time.Nanosecond); !errors.Is(err, cache.ErrInvalidLockTTL) || token != nil { + t.Fatalf("NewLock sub-ms ttl token=%v err=%v, want ErrInvalidLockTTL", token, err) + } + + token, err := cache.NewLock(ctx, "m10:ttl-valid", time.Second) + if err != nil || token == nil { + t.Fatalf("NewLock valid setup token=%v err=%v", token, err) + } + defer cache.Unlock(ctx, token) + + if err := cache.ExtendLock(ctx, token, 0); !errors.Is(err, cache.ErrInvalidLockTTL) { + t.Errorf("ExtendLock zero ttl err = %v, want ErrInvalidLockTTL", err) + } +} + +func TestWithLockAutoExtendRejectsNonPositiveInterval(t *testing.T) { + setupMiniRedis(t) + called := false + + err := cache.WithLockAutoExtend(context.Background(), "m10:interval", time.Second, 0, func() error { + called = true + return nil + }) + if !errors.Is(err, cache.ErrInvalidLockRetryInterval) { + t.Fatalf("WithLockAutoExtend zero interval err = %v, want ErrInvalidLockRetryInterval", err) + } + if called { + t.Fatal("WithLockAutoExtend should not run fn with invalid interval") + } +} + +func TestWithLockHeldReturnsErrLockNotAcquired(t *testing.T) { + setupMiniRedis(t) + ctx := context.Background() + key := "m10:not-acquired" + + holder, err := cache.NewLock(ctx, key, time.Second) + if err != nil || holder == nil { + t.Fatalf("setup holder token=%v err=%v", holder, err) + } + defer cache.Unlock(ctx, holder) + + called := false + err = cache.WithLock(ctx, key, time.Second, func() error { + called = true + return nil + }) + if !errors.Is(err, cache.ErrLockNotAcquired) { + t.Fatalf("WithLock held err = %v, want ErrLockNotAcquired", err) + } + if called { + t.Fatal("WithLock should not run fn when lock is held") + } +} + +func TestTryLockExhaustedReturnsErrLockNotAcquired(t *testing.T) { + setupMiniRedis(t) + ctx := context.Background() + key := "m10:try-exhausted" + + holder, err := cache.NewLock(ctx, key, time.Second) + if err != nil || holder == nil { + t.Fatalf("setup holder token=%v err=%v", holder, err) + } + defer cache.Unlock(ctx, holder) + + token, err := cache.TryLock(ctx, key, time.Second, time.Millisecond, 2) + if !errors.Is(err, cache.ErrLockNotAcquired) { + t.Fatalf("TryLock exhausted err = %v, want ErrLockNotAcquired", err) + } + if token != nil { + t.Fatalf("TryLock exhausted token = %v, want nil", token) + } +} diff --git a/cache/lock_test.go b/cache/lock_test.go index d257adf..c4c1411 100644 --- a/cache/lock_test.go +++ b/cache/lock_test.go @@ -51,6 +51,10 @@ func TestLockErrors(t *testing.T) { if ttl != 0 { t.Error("GetLockTTL should return 0 without Redis") } + + if err := cache.UnlockByKey(ctx, "test_key"); !errors.Is(err, cache.ErrRedisNotReady) { + t.Errorf("UnlockByKey without Redis should return ErrRedisNotReady, got %v", err) + } } func TestIncrDecr(t *testing.T) { @@ -146,4 +150,4 @@ func TestKFunctions(t *testing.T) { if sessionKey == "" { t.Error("KSession should not return empty string") } -} \ No newline at end of file +} diff --git a/compress/compress.go b/compress/compress.go index aea61c9..f78edaf 100644 --- a/compress/compress.go +++ b/compress/compress.go @@ -298,7 +298,7 @@ func UnzipWithOptions(zipPath, dstDir string, opts DecompressOptions) error { // 用绝对路径作目标锚定根，避免相对路径 + `..` 组合绕过前缀校验（C5a）。 absDst, err := filepath.Abs(dstDir) if err != nil { - absDst = dstDir + return fmt.Errorf("resolve unzip destination failed: %w", err) } absDst = filepath.Clean(absDst) diff --git a/jwt/jwt.go b/jwt/jwt.go index d2f1efe..12faaf3 100644 --- a/jwt/jwt.go +++ b/jwt/jwt.go @@ -53,6 +53,10 @@ var ( // 本实现仅支持 HMAC 族（HS256/HS384/HS512）；RS256 等非对称算法暂不支持， // 不再静默回退 HS256——避免用户误以为在用非对称算法、实则 HMAC，并助长算法混淆攻击。 ErrUnsupportedAlgorithm = errors.New("jwt: 不支持的签名算法（仅支持 HS256/HS384/HS512）") + // ErrInvalidExpiry 过期时间非法。签发非正过期时间的 token 会立即失效或产生不可预期会话语义。 + ErrInvalidExpiry = errors.New("jwt: 过期时间必须大于 0") + // ErrEmptyJTI 空 JTI 无法匹配任何正常 token，却会写入 jwt_bl: 这种永不命中的黑名单键。 + ErrEmptyJTI = errors.New("jwt: jti 不能为空") ) // validMethods 允许的签名算法名（HMAC 族）。ParseWithClaims 传 jwt.WithValidMethods 固定算法， @@ -125,6 +129,10 @@ func blacklistCtx() (context.Context, context.CancelFunc) { // 无 Redis 时返回 ErrBlacklistUnavailable（C9a 修复）：让调用方（RefreshToken/InvalidateToken） // 感知黑名单不可用并 fail-closed，避免无 Redis 时静默成功致撤销失效。 func (tb *TokenBlacklist) Add(jti string, expiry time.Time) error { + if strings.TrimSpace(jti) == "" { + return ErrEmptyJTI + } + client := tb.redisClient() if client == nil { // Redis 未启用，黑名单不可用——fail-closed 让调用方决策。 @@ -230,7 +238,17 @@ func (m *Manager) Blacklist() *TokenBlacklist { // GenerateToken 生成 JWT Token func GenerateToken(userID uint, username, role, userType string) (string, error) { cfg := config.Get() + var expiry time.Duration + if cfg != nil { + expiry = cfg.JWT.Expire + } + return generateTokenWithExpiry(cfg, userID, username, role, userType, expiry) +} +func generateTokenWithExpiry(cfg *config.Config, userID uint, username, role, userType string, expiry time.Duration) (string, error) { + if expiry <= 0 { + return "", ErrInvalidExpiry + } // P0：先校验密钥非空与算法受支持（fail-closed）。secretKey 亦守卫 cfg==nil， // 通过后 cfg 保证非空，后续访问 cfg.JWT.* 安全。 key, err := secretKey(cfg) @@ -255,7 +273,7 @@ func GenerateToken(userID uint, username, role, userType string) (string, error) UserType: userType, JTI: jti, RegisteredClaims: jwt.RegisteredClaims{ - ExpiresAt: jwt.NewNumericDate(time.Now().Add(cfg.JWT.Expire)), + ExpiresAt: jwt.NewNumericDate(time.Now().Add(expiry)), IssuedAt: jwt.NewNumericDate(time.Now()), NotBefore: jwt.NewNumericDate(time.Now()), Issuer: issuerOrDefault(cfg.JWT.Issuer), @@ -270,39 +288,10 @@ func GenerateToken(userID uint, username, role, userType string) (string, error) // GenerateTokenWithCustomExpiry 生成带自定义过期时间的 Token func GenerateTokenWithCustomExpiry(userID uint, username, role, userType string, expireSeconds int) (string, error) { cfg := config.Get() - - // P0：先校验密钥非空与算法受支持（fail-closed，见 GenerateToken）。 - key, err := secretKey(cfg) - if err != nil { - return "", err + if expireSeconds <= 0 { + return "", ErrInvalidExpiry } - method, err := signingMethod(cfg.JWT.Algorithm) - if err != nil { - return "", err - } - - jti, err := generateJTI() - if err != nil { - return "", err - } - - claims := Claims{ - UserID: userID, - Username: username, - Role: role, - UserType: userType, - JTI: jti, - RegisteredClaims: jwt.RegisteredClaims{ - ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Duration(expireSeconds) * time.Second)), - IssuedAt: jwt.NewNumericDate(time.Now()), - NotBefore: jwt.NewNumericDate(time.Now()), - Issuer: issuerOrDefault(cfg.JWT.Issuer), - ID: jti, - }, - } - - token := jwt.NewWithClaims(method, claims) - return token.SignedString(key) + return generateTokenWithExpiry(cfg, userID, username, role, userType, time.Duration(expireSeconds)*time.Second) } // issuerOrDefault 返回配置的 issuer，未配置时回退 "xlgo"。 @@ -329,11 +318,33 @@ func signingMethod(algorithm string) (jwt.SigningMethod, error) { } } +func parseOptions(cfg *config.Config) []jwt.ParserOption { + issuer := "xlgo" + if cfg != nil { + issuer = issuerOrDefault(cfg.JWT.Issuer) + } + return []jwt.ParserOption{ + jwt.WithValidMethods(validMethods), + jwt.WithIssuer(issuer), + } +} + +func validateIssuer(cfg *config.Config, claims *Claims) error { + want := "xlgo" + if cfg != nil { + want = issuerOrDefault(cfg.JWT.Issuer) + } + if claims == nil || claims.Issuer != want { + return ErrTokenInvalid + } + return nil +} + // ParseToken 解析 JWT Token func ParseToken(tokenString string) (*Claims, error) { cfg := config.Get() - token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), jwt.WithValidMethods(validMethods)) + token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), parseOptions(cfg)...) if err != nil { if errors.Is(err, jwt.ErrTokenExpired) { @@ -363,15 +374,22 @@ func ParseToken(tokenString string) (*Claims, error) { func InvalidateToken(tokenString string) error { cfg := config.Get() - token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), jwt.WithValidMethods(validMethods)) + opts := append(parseOptions(cfg), jwt.WithoutClaimsValidation()) + token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), opts...) if err != nil { - // Token 无效或已过期，无需加入黑名单 + // Token 签名/格式/issuer 无效，无需加入黑名单。 return nil } if claims, ok := token.Claims.(*Claims); ok { - if claims.JTI != "" && claims.ExpiresAt != nil { + if err := validateIssuer(cfg, claims); err != nil { + return nil + } + if strings.TrimSpace(claims.JTI) == "" { + return ErrEmptyJTI + } + if claims.ExpiresAt != nil { return currentBlacklist().Add(claims.JTI, claims.ExpiresAt.Time) } } @@ -382,6 +400,9 @@ func InvalidateToken(tokenString string) error { // InvalidateTokenByID 直接通过 JTI 使 Token 失效 // 参数: jti JWT ID，expiry 过期时间 func InvalidateTokenByID(jti string, expiry time.Time) error { + if strings.TrimSpace(jti) == "" { + return ErrEmptyJTI + } return currentBlacklist().Add(jti, expiry) } @@ -397,13 +418,24 @@ func RefreshToken(tokenString string) (string, error) { } // 将旧 Token 加入黑名单；失败则不签发新 token（C9b：禁止吞 Add 错误）。 - if claims.JTI != "" && claims.ExpiresAt != nil { + if strings.TrimSpace(claims.JTI) == "" { + return "", ErrEmptyJTI + } + if claims.ExpiresAt != nil { if err := currentBlacklist().Add(claims.JTI, claims.ExpiresAt.Time); err != nil { return "", fmt.Errorf("刷新令牌失败：旧令牌撤销失败: %w", err) } } - return GenerateToken(claims.UserID, claims.Username, claims.Role, claims.UserType) + cfg := config.Get() + var expiry time.Duration + if cfg != nil { + expiry = cfg.JWT.RefreshExpire + } + if expiry <= 0 && cfg != nil { + expiry = cfg.JWT.Expire + } + return generateTokenWithExpiry(cfg, claims.UserID, claims.Username, claims.Role, claims.UserType, expiry) } // GetJTI 从 Token 中提取 JTI（不验证签名） @@ -431,16 +463,19 @@ func IsTokenRevoked(jti string) bool { func GetClaimsFromToken(tokenString string) (*Claims, error) { cfg := config.Get() - token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), - jwt.WithValidMethods(validMethods), jwt.WithoutClaimsValidation()) + opts := append(parseOptions(cfg), jwt.WithoutClaimsValidation()) + token, err := jwt.ParseWithClaims(tokenString, &Claims{}, hmacKeyfunc(cfg), opts...) if err != nil { return nil, err } if claims, ok := token.Claims.(*Claims); ok { + if err := validateIssuer(cfg, claims); err != nil { + return nil, err + } return claims, nil } return nil, ErrTokenInvalid -} \ No newline at end of file +} diff --git a/jwt/jwt_test.go b/jwt/jwt_test.go index 8302700..6c2c96d 100644 --- a/jwt/jwt_test.go +++ b/jwt/jwt_test.go @@ -17,7 +17,7 @@ func setupTestConfig() { cfg := &config.Config{ JWT: config.JWTConfig{ Secret: "test-secret-key-1234567890123456789012", // ≥32 字节 - Expire: time.Hour, // 1小时 + Expire: time.Hour, // 1小时 }, } config.Set(cfg) @@ -384,3 +384,130 @@ func TestSupportedAlgorithmHS384(t *testing.T) { } } +func TestParseTokenRejectsWrongIssuer(t *testing.T) { + config.Set(&config.Config{JWT: config.JWTConfig{ + Secret: "test-secret-key-1234567890123456789012", + Expire: time.Hour, + Issuer: "trusted-issuer", + }}) + t.Cleanup(setupTestConfig) + + claims := jwt.Claims{ + UserID: 1, + Username: "attacker", + Role: "admin", + UserType: "super_admin", + RegisteredClaims: gojwt.RegisteredClaims{ + ExpiresAt: gojwt.NewNumericDate(time.Now().Add(time.Hour)), + IssuedAt: gojwt.NewNumericDate(time.Now()), + NotBefore: gojwt.NewNumericDate(time.Now()), + Issuer: "evil-issuer", + }, + } + token := signTokenForTest(t, claims) + + if _, err := jwt.ParseToken(token); !errors.Is(err, jwt.ErrTokenInvalid) { + t.Errorf("ParseToken wrong issuer err = %v, want ErrTokenInvalid", err) + } + if _, err := jwt.GetClaimsFromToken(token); !errors.Is(err, jwt.ErrTokenInvalid) { + t.Errorf("GetClaimsFromToken wrong issuer err = %v, want ErrTokenInvalid", err) + } +} + +func TestRefreshTokenUsesRefreshExpire(t *testing.T) { + config.Set(&config.Config{JWT: config.JWTConfig{ + Secret: "test-secret-key-1234567890123456789012", + Expire: time.Hour, + RefreshExpire: 4 * time.Hour, + }}) + t.Cleanup(setupTestConfig) + setupMiniRedis(t) + + token, err := jwt.GenerateToken(1, "testuser", "admin", "super_admin") + if err != nil { + t.Fatalf("GenerateToken: %v", err) + } + newToken, err := jwt.RefreshToken(token) + if err != nil { + t.Fatalf("RefreshToken: %v", err) + } + claims, err := jwt.ParseToken(newToken) + if err != nil { + t.Fatalf("ParseToken refreshed token: %v", err) + } + ttl := time.Until(claims.ExpiresAt.Time) + if ttl < 3*time.Hour || ttl > 4*time.Hour+time.Minute { + t.Errorf("refreshed token ttl = %v, want about 4h", ttl) + } +} + +func TestGenerateTokenWithCustomExpiryRejectsNonPositive(t *testing.T) { + setupTestConfig() + + for _, seconds := range []int{0, -1} { + if _, err := jwt.GenerateTokenWithCustomExpiry(1, "u", "admin", "admin", seconds); !errors.Is(err, jwt.ErrInvalidExpiry) { + t.Errorf("GenerateTokenWithCustomExpiry(%d) err = %v, want ErrInvalidExpiry", seconds, err) + } + } +} + +func TestInvalidateTokenByIDRejectsEmptyJTI(t *testing.T) { + setupTestConfig() + + if err := jwt.InvalidateTokenByID("", time.Now().Add(time.Hour)); !errors.Is(err, jwt.ErrEmptyJTI) { + t.Errorf("InvalidateTokenByID empty err = %v, want ErrEmptyJTI", err) + } +} + +func TestRefreshTokenRejectsEmptyJTI(t *testing.T) { + setupTestConfig() + setupMiniRedis(t) + + token := signTokenForTest(t, jwt.Claims{ + UserID: 1, + Username: "legacy", + Role: "admin", + UserType: "admin", + RegisteredClaims: gojwt.RegisteredClaims{ + ExpiresAt: gojwt.NewNumericDate(time.Now().Add(time.Hour)), + IssuedAt: gojwt.NewNumericDate(time.Now()), + NotBefore: gojwt.NewNumericDate(time.Now()), + Issuer: "xlgo", + }, + }) + + if _, err := jwt.RefreshToken(token); !errors.Is(err, jwt.ErrEmptyJTI) { + t.Errorf("RefreshToken empty JTI err = %v, want ErrEmptyJTI", err) + } +} + +func TestInvalidateTokenRejectsEmptyJTI(t *testing.T) { + setupTestConfig() + setupMiniRedis(t) + + token := signTokenForTest(t, jwt.Claims{ + UserID: 1, + Username: "legacy", + Role: "admin", + UserType: "admin", + RegisteredClaims: gojwt.RegisteredClaims{ + ExpiresAt: gojwt.NewNumericDate(time.Now().Add(time.Hour)), + IssuedAt: gojwt.NewNumericDate(time.Now()), + NotBefore: gojwt.NewNumericDate(time.Now()), + Issuer: "xlgo", + }, + }) + + if err := jwt.InvalidateToken(token); !errors.Is(err, jwt.ErrEmptyJTI) { + t.Errorf("InvalidateToken empty JTI err = %v, want ErrEmptyJTI", err) + } +} + +func signTokenForTest(t *testing.T, claims jwt.Claims) string { + t.Helper() + token, err := gojwt.NewWithClaims(gojwt.SigningMethodHS256, claims).SignedString([]byte("test-secret-key-1234567890123456789012")) + if err != nil { + t.Fatalf("sign token: %v", err) + } + return token +} diff --git a/sse/sse.go b/sse/sse.go index 38fdc95..9b868c5 100644 --- a/sse/sse.go +++ b/sse/sse.go @@ -14,12 +14,16 @@ package sse import ( "context" "encoding/json" + "errors" "fmt" "net/http" + "strings" "github.com/gin-gonic/gin" ) +var ErrInvalidEventName = errors.New("sse: event name must not contain CR or LF") + // SSEWriter SSE 写入器 type SSEWriter struct { writer gin.ResponseWriter @@ -54,10 +58,13 @@ func NewSSEWriter(c *gin.Context) (*SSEWriter, error) { // 写错误向上传播（C3b 修复）：旧实现丢弃 fmt.Fprintf 错误且恒 return nil， // 导致客户端断连后 Stream 守卫永不触发、消费循环不退出、上游 LLM 流持续运行。 func (w *SSEWriter) WriteEvent(event, data string) error { + if strings.ContainsAny(event, "\r\n") { + return ErrInvalidEventName + } if _, err := fmt.Fprintf(w.writer, "event: %s\n", event); err != nil { return err } - if _, err := fmt.Fprintf(w.writer, "data: %s\n\n", data); err != nil { + if err := w.writeDataLines(data); err != nil { return err } w.flusher.Flush() @@ -67,13 +74,25 @@ func (w *SSEWriter) WriteEvent(event, data string) error { // WriteMessage 写入消息（无事件类型） // 格式: data: \n\n func (w *SSEWriter) WriteMessage(data string) error { - if _, err := fmt.Fprintf(w.writer, "data: %s\n\n", data); err != nil { + if err := w.writeDataLines(data); err != nil { return err } w.flusher.Flush() return nil } +func (w *SSEWriter) writeDataLines(data string) error { + data = strings.ReplaceAll(data, "\r\n", "\n") + data = strings.ReplaceAll(data, "\r", "\n") + for _, line := range strings.Split(data, "\n") { + if _, err := fmt.Fprintf(w.writer, "data: %s\n", line); err != nil { + return err + } + } + _, err := fmt.Fprint(w.writer, "\n") + return err +} + // WriteJSON 写入 JSON 数据 func (w *SSEWriter) WriteJSON(event string, data any) error { jsonData, err := json.Marshal(data) diff --git a/sse/sse_test.go b/sse/sse_test.go index 434cdb7..313d02a 100644 --- a/sse/sse_test.go +++ b/sse/sse_test.go @@ -1,6 +1,7 @@ package sse_test import ( + "errors" "net/http/httptest" "strings" "testing" @@ -139,4 +140,48 @@ func TestSSEHeaders(t *testing.T) { if w.Header().Get("Connection") != "keep-alive" { t.Error("Connection should be keep-alive") } -} \ No newline at end of file +} + +func TestSSEWriterRejectsEventNewlineInjection(t *testing.T) { + r := setupTestRouter() + var writeErr error + r.GET("/sse", func(c *gin.Context) { + writer, _ := sse.NewSSEWriter(c) + writeErr = writer.WriteEvent("message\nevent: hacked", "safe") + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/sse", nil) + r.ServeHTTP(w, req) + + if !errors.Is(writeErr, sse.ErrInvalidEventName) { + t.Fatalf("WriteEvent newline err = %v, want ErrInvalidEventName", writeErr) + } + if strings.Contains(w.Body.String(), "event: hacked") { + t.Fatalf("response contains injected event: %q", w.Body.String()) + } +} + +func TestSSEWriterSplitsMultilineData(t *testing.T) { + r := setupTestRouter() + r.GET("/sse", func(c *gin.Context) { + writer, _ := sse.NewSSEWriter(c) + if err := writer.WriteMessage("first\nevent: injected\r\nsecond"); err != nil { + t.Errorf("WriteMessage: %v", err) + } + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/sse", nil) + r.ServeHTTP(w, req) + + body := w.Body.String() + if strings.Contains(body, "\nevent: injected") { + t.Fatalf("data newline escaped into event field: %q", body) + } + for _, want := range []string{"data: first", "data: event: injected", "data: second"} { + if !strings.Contains(body, want) { + t.Fatalf("body = %q, want %q", body, want) + } + } +} diff --git a/storage/storage.go b/storage/storage.go index d1aa108..96e00ca 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -244,7 +244,7 @@ func (s *LocalStorage) safeJoin(parts ...string) (string, error) { } // Upload 上传文件 -func (s *LocalStorage) Upload(file *multipart.FileHeader, subdir string) (string, error) { +func (s *LocalStorage) Upload(file *multipart.FileHeader, subdir string) (ret string, err error) { // 上传安全策略校验（大小 / 扩展名）；file.Size 由 multipart 解析时填，无需打开文件（C4b）。 if err := validateUploadSize(s.policy, file.Size); err != nil { return "", err @@ -294,13 +294,18 @@ func (s *LocalStorage) Upload(file *multipart.FileHeader, subdir string) (string if err != nil { return "", fmt.Errorf("创建文件失败: %w", err) } - defer dstFile.Close() + defer func() { + if cerr := dstFile.Close(); cerr != nil { + err = errors.Join(err, cerr) + } + if err != nil { + ret = "" + _ = os.Remove(dst) + } + }() // 复制文件内容。按策略实测封顶（P0）：超限即报错并清理已落盘的部分文件。 if _, err := io.Copy(dstFile, enforceUploadSize(s.policy, srcReader)); err != nil { - // 先关闭再删除（Windows 下删除打开中的文件会失败）；defer 的 Close 随后成为无害 no-op。 - _ = dstFile.Close() - _ = os.Remove(dst) return "", fmt.Errorf("保存文件失败: %w", err) } @@ -314,7 +319,7 @@ func (s *LocalStorage) Upload(file *multipart.FileHeader, subdir string) (string } // UploadFromBytes 从字节数组上传文件 -func (s *LocalStorage) UploadFromBytes(data []byte, filename, subdir string) (string, error) { +func (s *LocalStorage) UploadFromBytes(data []byte, filename, subdir string) (ret string, err error) { // 上传安全策略校验（C4b） if err := validateUploadSize(s.policy, int64(len(data))); err != nil { return "", err @@ -360,7 +365,15 @@ func (s *LocalStorage) UploadFromBytes(data []byte, filename, subdir string) (st if err != nil { return "", fmt.Errorf("创建文件失败: %w", err) } - defer dstFile.Close() + defer func() { + if cerr := dstFile.Close(); cerr != nil { + err = errors.Join(err, cerr) + } + if err != nil { + ret = "" + _ = os.Remove(dst) + } + }() // 写入文件内容 if _, err := io.Copy(dstFile, srcReader); err != nil { @@ -571,7 +584,11 @@ func (s *OSSStorage) GetURL(path string) string { // GetSignedURL 获取带签名的临时访问 URL（用于私有文件） func (s *OSSStorage) GetSignedURL(path string, expire time.Duration) (string, error) { - return s.bucket.SignURL(path, oss.HTTPGet, int64(expire.Seconds())) + key, err := sanitizeObjectKey(path) + if err != nil { + return "", err + } + return s.bucket.SignURL(key, oss.HTTPGet, int64(expire.Seconds())) } // Delete 删除 OSS 文件