[优化] POW 系统接口方案

This commit is contained in:
ryan
2026-06-06 12:13:55 +08:00
parent 4be7c19798
commit 6bf5023af1
20 changed files with 896 additions and 36 deletions
+221
View File
@@ -0,0 +1,221 @@
package cap
import (
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"strconv"
"strings"
"time"
)
const jwtHeaderB64 = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"
// ChallengeConfig holds parameters for the PoW challenge
type ChallengeConfig struct {
Count int // Number of puzzles (c)
Size int // Salt length (s)
Difficulty int // Difficulty prefix length (d)
ExpiresMs time.Duration // Challenge TTL
}
// ChallengeResponse is returned to the client
type ChallengeResponse struct {
Challenge struct {
C int `json:"c"`
S int `json:"s"`
D int `json:"d"`
} `json:"challenge"`
Token string `json:"token"`
Expires int64 `json:"expires"` // ms timestamp
}
// ChallengePayload represents the signed JWT payload
type ChallengePayload struct {
Nonce string `json:"n"`
Count int `json:"c"`
Size int `json:"s"`
Difficulty int `json:"d"`
Expires int64 `json:"exp"` // ms timestamp
IssuedAt int64 `json:"iat"` // ms timestamp
Scope string `json:"sk,omitempty"`
}
// RedeemRequest payload sent by client
type RedeemRequest struct {
Token string `json:"token"`
Solutions []int `json:"solutions"`
}
// RedeemResponse returned to client after verification
type RedeemResponse struct {
Success bool `json:"success"`
Token string `json:"token,omitempty"`
Expires int64 `json:"expires,omitempty"`
Error string `json:"error,omitempty"`
}
func b64urlEncode(data []byte) string {
return base64.RawURLEncoding.EncodeToString(data)
}
func b64urlDecode(str string) ([]byte, error) {
return base64.RawURLEncoding.DecodeString(str)
}
func randomHex(byteLen int) string {
bytes := make([]byte, byteLen)
if _, err := rand.Read(bytes); err != nil {
panic(err)
}
return hex.EncodeToString(bytes)
}
func jwtSign(payload []byte, secret []byte) string {
body := b64urlEncode(payload)
sigInput := jwtHeaderB64 + "." + body
mac := hmac.New(sha256.New, secret)
mac.Write([]byte(sigInput))
sig := mac.Sum(nil)
return sigInput + "." + b64urlEncode(sig)
}
func jwtVerify(token string, secret []byte) ([]byte, error) {
parts := strings.Split(token, ".")
if len(parts) != 3 {
return nil, errors.New("invalid token format")
}
if parts[0] != jwtHeaderB64 {
return nil, errors.New("invalid header")
}
sigInput := parts[0] + "." + parts[1]
mac := hmac.New(sha256.New, secret)
mac.Write([]byte(sigInput))
expectedSig := mac.Sum(nil)
actualSig, err := b64urlDecode(parts[2])
if err != nil {
return nil, err
}
if !hmac.Equal(expectedSig, actualSig) {
return nil, errors.New("signature mismatch")
}
payload, err := b64urlDecode(parts[1])
if err != nil {
return nil, err
}
return payload, nil
}
func jwtSigHex(token string) string {
parts := strings.Split(token, ".")
if len(parts) != 3 {
return ""
}
sigBytes, err := b64urlDecode(parts[2])
if err != nil {
return ""
}
return hex.EncodeToString(sigBytes)
}
// GenerateChallenge produces a new challenge and signed token
func GenerateChallenge(secret []byte, conf ChallengeConfig, scope string) (*ChallengeResponse, error) {
if conf.Count <= 0 {
conf.Count = 50
}
if conf.Size <= 0 {
conf.Size = 32
}
if conf.Difficulty <= 0 {
conf.Difficulty = 4
}
if conf.ExpiresMs <= 0 {
conf.ExpiresMs = 10 * time.Minute
}
now := time.Now().UnixNano() / int64(time.Millisecond)
expires := now + int64(conf.ExpiresMs/time.Millisecond)
payload := ChallengePayload{
Nonce: randomHex(25),
Count: conf.Count,
Size: conf.Size,
Difficulty: conf.Difficulty,
Expires: expires,
IssuedAt: now,
Scope: scope,
}
payloadBytes, err := json.Marshal(payload)
if err != nil {
return nil, err
}
token := jwtSign(payloadBytes, secret)
resp := &ChallengeResponse{
Token: token,
Expires: expires,
}
resp.Challenge.C = conf.Count
resp.Challenge.S = conf.Size
resp.Challenge.D = conf.Difficulty
return resp, nil
}
// VerifyChallengeSolutions verifies client submitted solutions
func VerifyChallengeSolutions(token string, solutions []int, secret []byte, expectedScope string) (*ChallengePayload, error) {
payloadBytes, err := jwtVerify(token, secret)
if err != nil {
return nil, errors.New("invalid_token")
}
var payload ChallengePayload
if err := json.Unmarshal(payloadBytes, &payload); err != nil {
return nil, errors.New("invalid_token")
}
if expectedScope != "" && payload.Scope != expectedScope {
return nil, errors.New("scope_mismatch")
}
now := time.Now().UnixNano() / int64(time.Millisecond)
if payload.Expires < now {
return nil, errors.New("expired")
}
if len(solutions) != payload.Count {
return nil, errors.New("invalid_solutions")
}
tokenFnv := fnv1a(token)
for i := 0; i < payload.Count; i++ {
idxStr := strconv.Itoa(i + 1)
saltSeed := fnv1aResume(tokenFnv, idxStr)
targetSeed := fnv1aResume(saltSeed, "d")
salt := prngFromHash(saltSeed, payload.Size)
target := prngFromHash(targetSeed, payload.Difficulty)
hashInput := salt + strconv.Itoa(solutions[i])
hashBytes := sha256.Sum256([]byte(hashInput))
hashHex := hex.EncodeToString(hashBytes[:])
if !strings.HasPrefix(hashHex, target) {
return nil, errors.New("invalid_solution")
}
}
return &payload, nil
}
+93
View File
@@ -0,0 +1,93 @@
package cap
import (
"context"
"crypto/sha256"
"encoding/hex"
"strconv"
"strings"
"testing"
"time"
)
func TestCapFullFlow(t *testing.T) {
secret := []byte("a-very-long-secret-key-at-least-16-bytes")
store := 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"
resp, err := manager.Generate(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 := make([]int, resp.Challenge.C)
tokenFnv := fnv1a(resp.Token)
for i := 0; i < resp.Challenge.C; i++ {
idxStr := strconv.Itoa(i + 1)
saltSeed := fnv1aResume(tokenFnv, idxStr)
targetSeed := fnv1aResume(saltSeed, "d")
salt := prngFromHash(saltSeed, resp.Challenge.S)
target := prngFromHash(targetSeed, resp.Challenge.D)
// Brute force the PoW solution
var found bool
for nonce := 0; nonce < 1000000; nonce++ {
hashInput := salt + strconv.Itoa(nonce)
hashBytes := sha256.Sum256([]byte(hashInput))
hashHex := hex.EncodeToString(hashBytes[:])
if strings.HasPrefix(hashHex, target) {
solutions[i] = nonce
found = true
break
}
}
if !found {
t.Fatalf("Failed to solve puzzle %d", i)
}
}
// Redeem
ctx := context.Background()
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)")
}
}
+165
View File
@@ -0,0 +1,165 @@
package cap
import (
"context"
"crypto/sha256"
"encoding/hex"
"strconv"
"strings"
"time"
)
// 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 Store
}
// NewManager creates a new CAPTCHA Manager
func NewManager(conf Config, store Store) *Manager {
if conf.ChallengeCount <= 0 {
conf.ChallengeCount = 50
}
if conf.ChallengeSize <= 0 {
conf.ChallengeSize = 32
}
if conf.ChallengeDifficulty <= 0 {
conf.ChallengeDifficulty = 4
}
if conf.ChallengeTTL <= 0 {
conf.ChallengeTTL = 10 * time.Minute
}
if conf.TokenTTL <= 0 {
conf.TokenTTL = 20 * time.Minute
}
return &Manager{
conf: conf,
store: store,
}
}
// Generate creates a challenge response
func (m *Manager) Generate(scope string) (*ChallengeResponse, error) {
c := ChallengeConfig{
Count: m.conf.ChallengeCount,
Size: m.conf.ChallengeSize,
Difficulty: m.conf.ChallengeDifficulty,
ExpiresMs: m.conf.ChallengeTTL,
}
return GenerateChallenge(m.conf.Secret, c, scope)
}
// 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 := jwtSigHex(token)
if sigHex == "" {
return &RedeemResponse{Success: false, Error: "invalid_token"}, nil
}
// Replay prevention: check if this JWT signature has already been used
nonceKey := "cap:nonce:" + sigHex
_, exists, err := m.store.Get(ctx, nonceKey)
if err != nil {
return &RedeemResponse{Success: false, Error: "nonce_store_error"}, err
}
if exists {
return &RedeemResponse{Success: false, Error: "already_redeemed"}, nil
}
payload, err := VerifyChallengeSolutions(token, solutions, m.conf.Secret, scope)
if err != nil {
return &RedeemResponse{Success: false, Error: err.Error()}, nil
}
// Verification succeeded. Consume the nonce.
now := time.Now().UnixNano() / int64(time.Millisecond)
ttlMs := time.Duration(payload.Expires-now) * time.Millisecond
if ttlMs < time.Second {
ttlMs = time.Second
}
if err := m.store.Set(ctx, nonceKey, "1", ttlMs); err != nil {
return &RedeemResponse{Success: false, Error: "nonce_store_error"}, err
}
// Generate a redeem token formatted as "id:verToken"
id := randomHex(8)
verToken := randomHex(15)
verHashBytes := sha256.Sum256([]byte(verToken))
verHashHex := hex.EncodeToString(verHashBytes[:])
tokenKey := "cap:token:" + id + ":" + verHashHex
tokenExpires := time.Now().Add(m.conf.TokenTTL)
// Value stored is "expiresNano|scope"
storeVal := strconv.FormatInt(tokenExpires.UnixNano(), 10) + "|" + scope
if err := m.store.Set(ctx, tokenKey, storeVal, m.conf.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) != 2 {
return false, nil
}
id := parts[0]
verToken := parts[1]
verHashBytes := sha256.Sum256([]byte(verToken))
verHashHex := hex.EncodeToString(verHashBytes[:])
tokenKey := "cap:token:" + id + ":" + verHashHex
val, exists, err := m.store.Get(ctx, tokenKey)
if err != nil {
return false, err
}
if !exists {
return false, nil
}
// Single-use: consume/delete the token immediately
_ = m.store.Delete(ctx, tokenKey)
valParts := strings.Split(val, "|")
if len(valParts) != 2 {
return false, nil
}
expNano, err := strconv.ParseInt(valParts[0], 10, 64)
if err != nil {
return false, nil
}
tokenScope := valParts[1]
if expectedScope != "" && tokenScope != expectedScope {
return false, nil
}
if time.Now().UnixNano() > expNano {
return false, nil // Expired
}
return true, nil
}
+40
View File
@@ -0,0 +1,40 @@
package cap
import (
"net/http"
"github.com/gin-gonic/gin"
)
// VerifyMiddleware returns a Gin middleware that checks and consumes the X-Cap-Token header.
// enabledFunc is an optional callback allowing dynamic check of whether captcha protection is turned on.
func (m *Manager) VerifyMiddleware(scope string, enabledFunc func() bool) gin.HandlerFunc {
return func(c *gin.Context) {
if enabledFunc != nil && !enabledFunc() {
c.Next()
return
}
token := c.GetHeader("X-Cap-Token")
if token == "" {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"error": "验证码验证失败,缺少验证码凭证",
})
c.Abort()
return
}
valid, err := m.VerifyToken(c.Request.Context(), token, scope)
if err != nil || !valid {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"error": "验证码校验失败或已过期,请重试",
})
c.Abort()
return
}
c.Next()
}
}
+45
View File
@@ -0,0 +1,45 @@
package cap
import (
"fmt"
"strings"
)
// fnv1a returns the 32-bit FNV-1a hash of a string
func fnv1a(str string) uint32 {
var hash uint32 = 2166136261
for i := 0; i < len(str); i++ {
hash ^= uint32(str[i])
hash += (hash << 1) + (hash << 4) + (hash << 7) + (hash << 8) + (hash << 24)
}
return hash
}
// fnv1aResume resumes FNV-1a hashing from a given state
func fnv1aResume(state uint32, str string) uint32 {
h := state
for i := 0; i < len(str); i++ {
h ^= uint32(str[i])
h += (h << 1) + (h << 4) + (h << 7) + (h << 8) + (h << 24)
}
return h
}
// prng generates a hex string of specified length using a seed
func prng(seed string, length int) string {
return prngFromHash(fnv1a(seed), length)
}
// prngFromHash generates a hex string of specified length using an initial hash state
func prngFromHash(initialHash uint32, length int) string {
state := initialHash
var result strings.Builder
for result.Len() < length {
state ^= state << 13
state ^= state >> 17
state ^= state << 5
hexStr := fmt.Sprintf("%08x", state)
result.WriteString(hexStr)
}
return result.String()[:length]
}
+90
View File
@@ -0,0 +1,90 @@
package cap
import (
"context"
"sync"
"time"
)
// Store defines the storage interface for challenge nonces and verification tokens
type Store interface {
Get(ctx context.Context, key string) (string, bool, error)
Set(ctx context.Context, key string, val string, ttl time.Duration) error
Delete(ctx context.Context, key string) error
}
type memoryItem struct {
value string
expiresAt time.Time
}
// MemoryStore is a thread-safe in-memory implementation of Store
type MemoryStore struct {
items map[string]memoryItem
mu sync.RWMutex
}
// NewMemoryStore creates and initializes a new MemoryStore
func NewMemoryStore(cleanupInterval time.Duration) *MemoryStore {
store := &MemoryStore{
items: make(map[string]memoryItem),
}
if cleanupInterval > 0 {
go store.startCleanupLoop(cleanupInterval)
}
return store
}
func (s *MemoryStore) Get(ctx context.Context, key string) (string, bool, error) {
s.mu.RLock()
item, found := s.items[key]
s.mu.RUnlock()
if !found {
return "", false, nil
}
if time.Now().After(item.expiresAt) {
s.mu.Lock()
delete(s.items, key)
s.mu.Unlock()
return "", false, nil
}
return item.value, true, nil
}
func (s *MemoryStore) Set(ctx context.Context, key string, val string, ttl time.Duration) error {
s.mu.Lock()
defer s.mu.Unlock()
s.items[key] = memoryItem{
value: val,
expiresAt: time.Now().Add(ttl),
}
return nil
}
func (s *MemoryStore) Delete(ctx context.Context, key string) error {
s.mu.Lock()
defer s.mu.Unlock()
delete(s.items, key)
return nil
}
func (s *MemoryStore) startCleanupLoop(interval time.Duration) {
ticker := time.NewTicker(interval)
for range ticker.C {
s.cleanupExpired()
}
}
func (s *MemoryStore) cleanupExpired() {
now := time.Now()
s.mu.Lock()
defer s.mu.Unlock()
for k, v := range s.items {
if now.After(v.expiresAt) {
delete(s.items, k)
}
}
}