refactor(backend): extract to pkg/cap

2026-10-04 23:16:37 +08:00 · 2026-06-15 16:19:18 +08:00
parent 84ae4ec27e
commit b3ed94342c
63 changed files with 770 additions and 706 deletions
@@ -0,0 +1,282 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+// Package cap provides CAPTCHA and proof-of-work (PoW) verification services.
+package cap
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/Rain-kl/Wavelet/internal/config"
+	"github.com/Rain-kl/Wavelet/internal/db"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	pkgcap "github.com/Rain-kl/Wavelet/pkg/cap"
+)
+
+const (
+	managerDefaultChallengeCount = 1
+	managerDefaultChallengeSize  = 32
+	defaultChallengeDifficulty   = 4
+	defaultChallengeTTL          = 10 * time.Minute
+	defaultTokenTTL              = 20 * time.Minute
+	redeemTokenIDLength          = 8  // 兑换 Token ID 字节长度
+	redeemVerTokenLength         = 15 // 兑换验证 Token 字节长度
+	tokenPartsCount              = 2  // 兑换 Token 由两部分组成
+	valuePartsCount              = 2  // 存储值由 scope 和过期时间组成
+)
+
+// Config holds settings for the CAPTCHA manager
+type Config struct {
+	Secret              []byte        // HMAC signing key
+	ChallengeCount      int           // Number of PoW puzzles
+	ChallengeSize       int           // Size of the salt string
+	ChallengeDifficulty int           // Length of difficulty target prefix
+	ChallengeTTL        time.Duration // Lifespan of the challenge JWT
+	TokenTTL            time.Duration // Lifespan of the redeem token
+}
+
+// Manager orchestrates challenge generation and solution validation
+type Manager struct {
+	conf  Config
+	store pkgcap.Store
+}
+
+// NewManager creates a new CAPTCHA Manager
+func NewManager(conf Config, store pkgcap.Store) *Manager {
+	if conf.ChallengeCount <= 0 {
+		conf.ChallengeCount = managerDefaultChallengeCount
+	}
+	if conf.ChallengeSize <= 0 {
+		conf.ChallengeSize = managerDefaultChallengeSize
+	}
+	if conf.ChallengeDifficulty <= 0 {
+		conf.ChallengeDifficulty = defaultChallengeDifficulty
+	}
+	if conf.ChallengeTTL <= 0 {
+		conf.ChallengeTTL = defaultChallengeTTL
+	}
+	if conf.TokenTTL <= 0 {
+		conf.TokenTTL = defaultTokenTTL
+	}
+	return &Manager{
+		conf:  conf,
+		store: store,
+	}
+}
+
+// Generate creates a challenge response
+func (m *Manager) Generate(ctx context.Context, scope string) (*pkgcap.ChallengeResponse, error) {
+	c := pkgcap.ChallengeConfig{
+		Count:      m.getChallengeCount(ctx),
+		Size:       m.getChallengeSize(ctx),
+		Difficulty: m.getChallengeDifficulty(ctx),
+		Expires:    m.getChallengeTTL(ctx),
+	}
+	return pkgcap.GenerateChallenge(m.conf.Secret, c, scope)
+}
+
+// RedeemResponse is returned to the client on redeem
+type RedeemResponse struct {
+	Success bool   `json:"success"`
+	Token   string `json:"token,omitempty"`
+	Expires int64  `json:"expires,omitempty"`
+	Error   string `json:"error,omitempty"`
+}
+
+// Redeem verifies PoW solutions and returns a one-time redeem token
+func (m *Manager) Redeem(ctx context.Context, token string, solutions []int, scope string) (*RedeemResponse, error) {
+	sigHex := pkgcap.JwtSigHex(token)
+	if sigHex == "" {
+		return &RedeemResponse{Success: false, Error: "invalid_token"}, nil
+	}
+
+	nonceKey := "cap:nonce:" + sigHex
+
+	// Atomically claim the nonce slot BEFORE verifying solutions.
+	payload, err := pkgcap.VerifyChallengeSolutions(token, solutions, m.conf.Secret, scope)
+	if err != nil {
+		return &RedeemResponse{Success: false, Error: err.Error()}, nil //nolint:nilerr // expected behavior: validation error is returned as response, not system error
+	}
+
+	// Calculate remaining lifetime of the challenge JWT for the nonce TTL.
+	now := time.Now().UnixNano() / int64(time.Millisecond)
+	nonceTTL := time.Duration(payload.Expires-now) * time.Millisecond
+	if nonceTTL < time.Second {
+		nonceTTL = time.Second
+	}
+
+	// Atomic claim: if another goroutine already redeemed this JWT the SetNX
+	// will return false and we reject the request without issuing a token.
+	set, err := m.store.SetNX(ctx, nonceKey, "1", nonceTTL)
+	if err != nil {
+		return &RedeemResponse{Success: false, Error: "nonce_store_error"}, err
+	}
+	if !set {
+		return &RedeemResponse{Success: false, Error: "already_redeemed"}, nil
+	}
+
+	// Generate a redeem token formatted as "id:verToken"
+	id := pkgcap.RandomHex(redeemTokenIDLength)
+	verToken := pkgcap.RandomHex(redeemVerTokenLength)
+	verHashBytes := sha256.Sum256([]byte(verToken))
+	verHashHex := hex.EncodeToString(verHashBytes[:])
+
+	tokenKey := "cap:token:" + id + ":" + verHashHex
+	tokenTTL := m.getTokenTTL(ctx)
+	tokenExpires := time.Now().Add(tokenTTL)
+
+	// Value stored is "expiresNano|scope"
+	storeVal := strconv.FormatInt(tokenExpires.UnixNano(), 10) + "|" + scope
+
+	if err := m.store.Set(ctx, tokenKey, storeVal, tokenTTL); err != nil {
+		return &RedeemResponse{Success: false, Error: "token_store_error"}, err
+	}
+
+	return &RedeemResponse{
+		Success: true,
+		Token:   id + ":" + verToken,
+		Expires: tokenExpires.UnixNano() / int64(time.Millisecond),
+	}, nil
+}
+
+// VerifyToken validates and consumes the redeem token (single-use).
+func (m *Manager) VerifyToken(ctx context.Context, token string, expectedScope string) (bool, error) {
+	if token == "" {
+		return false, nil
+	}
+	parts := strings.Split(token, ":")
+	if len(parts) != tokenPartsCount {
+		return false, nil
+	}
+	id := parts[0]
+	verToken := parts[1]
+
+	verHashBytes := sha256.Sum256([]byte(verToken))
+	verHashHex := hex.EncodeToString(verHashBytes[:])
+
+	tokenKey := "cap:token:" + id + ":" + verHashHex
+
+	// Atomically retrieve-and-delete
+	val, exists, err := sGetAndDelete(ctx, m.store, tokenKey)
+	if err != nil {
+		return false, err
+	}
+	if !exists {
+		return false, nil
+	}
+
+	valParts := strings.Split(val, "|")
+	if len(valParts) != valuePartsCount {
+		return false, nil
+	}
+
+	expNano, err := strconv.ParseInt(valParts[0], 10, 64)
+	if err != nil {
+		return false, nil //nolint:nilerr // expected behavior: invalid format is treated as validation failure, not system error
+	}
+	tokenScope := valParts[1]
+
+	if expectedScope != "" && tokenScope != expectedScope {
+		return false, nil
+	}
+
+	if time.Now().UnixNano() > expNano {
+		return false, nil // Expired
+	}
+
+	return true, nil
+}
+
+// sGetAndDelete safely calls store.GetAndDelete, treating a nil store as a miss.
+func sGetAndDelete(ctx context.Context, store pkgcap.Store, key string) (string, bool, error) {
+	if store == nil {
+		return "", false, nil
+	}
+	return store.GetAndDelete(ctx, key)
+}
+
+func (m *Manager) getChallengeCount(ctx context.Context) int {
+	val, err := model.GetIntByKey(ctx, model.ConfigKeyCapChallengeCount)
+	if err != nil || val <= 0 {
+		return m.conf.ChallengeCount
+	}
+	return val
+}
+
+func (m *Manager) getChallengeSize(ctx context.Context) int {
+	val, err := model.GetIntByKey(ctx, model.ConfigKeyCapChallengeSize)
+	if err != nil || val <= 0 {
+		return m.conf.ChallengeSize
+	}
+	return val
+}
+
+func (m *Manager) getChallengeDifficulty(ctx context.Context) int {
+	val, err := model.GetIntByKey(ctx, model.ConfigKeyCapChallengeDifficulty)
+	if err != nil || val <= 0 {
+		return m.conf.ChallengeDifficulty
+	}
+	return val
+}
+
+func (m *Manager) getChallengeTTL(ctx context.Context) time.Duration {
+	val, err := model.GetIntByKey(ctx, model.ConfigKeyCapChallengeTTL)
+	if err != nil || val <= 0 {
+		return m.conf.ChallengeTTL
+	}
+	return time.Duration(val) * time.Second
+}
+
+func (m *Manager) getTokenTTL(ctx context.Context) time.Duration {
+	val, err := model.GetIntByKey(ctx, model.ConfigKeyCapTokenTTL)
+	if err != nil || val <= 0 {
+		return m.conf.TokenTTL
+	}
+	return time.Duration(val) * time.Second
+}
+
+var (
+	defaultManager *Manager
+	once           sync.Once
+)
+
+// GetDefaultManager yields the global singleton CAPTCHA manager
+func GetDefaultManager() *Manager {
+	once.Do(func() {
+		var secret []byte
+		if config.Config != nil && config.Config.App.SessionSecret != "" {
+			secret = []byte(config.Config.App.SessionSecret)
+		} else {
+			secret = []byte("default-captcha-secret-key-at-least-16-bytes")
+		}
+
+		challengeCount := managerDefaultChallengeCount
+		challengeSize := managerDefaultChallengeSize
+		challengeDifficulty := defaultChallengeDifficulty
+		challengeTTL := defaultChallengeTTL
+		tokenTTL := defaultTokenTTL
+
+		var store pkgcap.Store
+		if config.Config != nil && config.Config.Redis.Enabled && db.Redis != nil {
+			store = pkgcap.NewRedisStore(db.Redis)
+		} else {
+			store = pkgcap.NewMemoryStore(1 * time.Minute)
+		}
+
+		defaultManager = NewManager(Config{
+			Secret:              secret,
+			ChallengeCount:      challengeCount,
+			ChallengeSize:       challengeSize,
+			ChallengeDifficulty: challengeDifficulty,
+			ChallengeTTL:        challengeTTL,
+			TokenTTL:            tokenTTL,
+		}, store)
+	})
+	return defaultManager
+}
@@ -0,0 +1,172 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package cap
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	pkgcap "github.com/Rain-kl/Wavelet/pkg/cap"
+)
+
+func TestCapFullFlow(t *testing.T) {
+	secret := []byte("a-very-long-secret-key-at-least-16-bytes")
+	store := pkgcap.NewMemoryStore(1 * time.Minute)
+
+	manager := NewManager(Config{
+		Secret:              secret,
+		ChallengeCount:      3, // small count for fast test
+		ChallengeSize:       32,
+		ChallengeDifficulty: 3, // small difficulty for fast test
+		ChallengeTTL:        5 * time.Second,
+		TokenTTL:            10 * time.Second,
+	}, store)
+
+	scope := "test-scope"
+	ctx := context.Background()
+	resp, err := manager.Generate(ctx, scope)
+	if err != nil {
+		t.Fatalf("Generate failed: %v", err)
+	}
+
+	if resp.Challenge.C != 3 {
+		t.Errorf("Expected count 3, got %d", resp.Challenge.C)
+	}
+
+	// Solve the challenge (acting as client)
+	solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
+
+	// Redeem
+	redeemResp, err := manager.Redeem(ctx, resp.Token, solutions, scope)
+	if err != nil {
+		t.Fatalf("Redeem failed: %v", err)
+	}
+	if !redeemResp.Success {
+		t.Fatalf("Redeem returned success=false: %s", redeemResp.Error)
+	}
+	if redeemResp.Token == "" {
+		t.Fatalf("Expected token, got empty")
+	}
+
+	// Verify the token
+	valid, err := manager.VerifyToken(ctx, redeemResp.Token, scope)
+	if err != nil {
+		t.Fatalf("VerifyToken failed: %v", err)
+	}
+	if !valid {
+		t.Fatalf("Expected redeem token to be valid")
+	}
+
+	// Verify token is one-time use
+	validAgain, err := manager.VerifyToken(ctx, redeemResp.Token, scope)
+	if err != nil {
+		t.Fatalf("VerifyToken second call failed: %v", err)
+	}
+	if validAgain {
+		t.Fatalf("Expected redeem token to be single-use (invalidated after verification)")
+	}
+}
+
+// TestRedeemConcurrentRace verifies that when N goroutines simultaneously call
+// Redeem with the same challenge JWT, exactly one succeeds and the rest are
+// rejected with "already_redeemed". This guards against the TOCTOU fix.
+func TestRedeemConcurrentRace(t *testing.T) {
+	const goroutines = 50
+
+	secret := []byte("race-test-secret-key-at-least-16-bytes")
+	store := pkgcap.NewMemoryStore(1 * time.Minute)
+	manager := NewManager(Config{
+		Secret:              secret,
+		ChallengeCount:      1,
+		ChallengeSize:       32,
+		ChallengeDifficulty: 3,
+		ChallengeTTL:        30 * time.Second,
+		TokenTTL:            30 * time.Second,
+	}, store)
+
+	ctx := context.Background()
+	resp, err := manager.Generate(ctx, "login")
+	if err != nil {
+		t.Fatalf("Generate failed: %v", err)
+	}
+	solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
+
+	var (
+		wg      sync.WaitGroup
+		success atomic.Int32
+		barrier = make(chan struct{}) // synchronise goroutine start
+	)
+
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-barrier // wait for the gun
+			r, _ := manager.Redeem(ctx, resp.Token, solutions, "login")
+			if r != nil && r.Success {
+				success.Add(1)
+			}
+		}()
+	}
+	close(barrier) // fire all goroutines at once
+	wg.Wait()
+
+	if n := success.Load(); n != 1 {
+		t.Fatalf("Expected exactly 1 successful Redeem, got %d", n)
+	}
+}
+
+// TestVerifyTokenConcurrentRace verifies that when N goroutines simultaneously
+// call VerifyToken with the same cap token, exactly one succeeds and the rest
+// fail. This guards against the GetAndDelete fix.
+func TestVerifyTokenConcurrentRace(t *testing.T) {
+	const goroutines = 50
+
+	secret := []byte("race-test-secret-key-at-least-16-bytes")
+	store := pkgcap.NewMemoryStore(1 * time.Minute)
+	manager := NewManager(Config{
+		Secret:              secret,
+		ChallengeCount:      1,
+		ChallengeSize:       32,
+		ChallengeDifficulty: 3,
+		ChallengeTTL:        30 * time.Second,
+		TokenTTL:            30 * time.Second,
+	}, store)
+
+	ctx := context.Background()
+	resp, _ := manager.Generate(ctx, "login")
+	solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
+	redeemResp, err := manager.Redeem(ctx, resp.Token, solutions, "login")
+	if err != nil || !redeemResp.Success {
+		t.Fatalf("Redeem failed: %v %+v", err, redeemResp)
+	}
+	capToken := redeemResp.Token
+
+	var (
+		wg      sync.WaitGroup
+		success atomic.Int32
+		barrier = make(chan struct{})
+	)
+
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-barrier
+			ok, _ := manager.VerifyToken(ctx, capToken, "login")
+			if ok {
+				success.Add(1)
+			}
+		}()
+	}
+	close(barrier)
+	wg.Wait()
+
+	if n := success.Load(); n != 1 {
+		t.Fatalf("Expected exactly 1 successful VerifyToken, got %d", n)
+	}
+}
@@ -11,10 +11,10 @@ import (
 	"time"

 	"github.com/Rain-kl/Wavelet/internal/db"
-	"github.com/Rain-kl/Wavelet/internal/logger"
 	"github.com/Rain-kl/Wavelet/internal/model"
 	"github.com/Rain-kl/Wavelet/internal/storage"
 	"github.com/Rain-kl/Wavelet/internal/task"
+	"github.com/Rain-kl/Wavelet/pkg/logger"
 	"gorm.io/gorm"
 )