heicode-win/heicode/service/nonce_store.go

package service

import (
	"context"
	"sync"
	"time"

	"github.com/heicode/manager/common"
)

// NonceStore detects request replays by recording every (device_id, nonce)
// pair the device-signature middleware sees and refusing the second
// occurrence within a TTL window.
//
// Redis is the primary backend; when REDIS_CONN_STRING is unset we fall
// back to an in-memory sync.Map with a janitor goroutine. The fallback is
// fine for single-instance dev / SQLite deployments — replay protection
// for multi-instance production REQUIRES Redis (otherwise Pod A doesn't
// know what Pod B has seen).

type memoryNonceEntry struct {
	expiresAt time.Time
}

var (
	memoryNonces      sync.Map // key: string -> *memoryNonceEntry
	memoryJanitorOnce sync.Once
)

// startMemoryNonceJanitor sweeps expired entries every minute. Idempotent.
func startMemoryNonceJanitor() {
	memoryJanitorOnce.Do(func() {
		go func() {
			ticker := time.NewTicker(1 * time.Minute)
			defer ticker.Stop()
			for range ticker.C {
				now := time.Now()
				memoryNonces.Range(func(k, v any) bool {
					entry := v.(*memoryNonceEntry)
					if now.After(entry.expiresAt) {
						memoryNonces.Delete(k)
					}
					return true
				})
			}
		}()
	})
}

// MarkNonceUsed atomically records the (deviceId, nonce) pair as seen.
// Returns (ok=true) if this is the first time we've seen it within
// the TTL window. Returns (ok=false) if it was already used — the
// caller should reject the request as a replay.
//
// Never returns an error in normal operation; even if Redis is slow or
// failing it falls back to the memory store. A real error (programming
// bug) bubbles up and the caller can fail closed.
func MarkNonceUsed(deviceId, nonce string, ttl time.Duration) (bool, error) {
	key := "nonce:" + deviceId + ":" + nonce

	if common.RedisEnabled && common.RDB != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		// SET NX EX: returns true if the key was set (i.e. didn't exist).
		ok, err := common.RDB.SetNX(ctx, key, "1", ttl).Result()
		if err == nil {
			return ok, nil
		}
		// On Redis error, fall through to memory to fail soft. This is a
		// deliberate trade-off: a brief Redis outage briefly weakens
		// replay protection across instances, but doesn't kill the API.
		common.SysLog("nonce store: Redis SetNX failed, falling back to memory: " + err.Error())
	}

	startMemoryNonceJanitor()
	now := time.Now()
	expiresAt := now.Add(ttl)
	_, loaded := memoryNonces.LoadOrStore(key, &memoryNonceEntry{expiresAt: expiresAt})
	if loaded {
		// Key existed. Check whether it was a stale entry the janitor hasn't
		// reaped yet — if so, overwrite and treat as fresh.
		if entry, ok := memoryNonces.Load(key); ok {
			if now.After(entry.(*memoryNonceEntry).expiresAt) {
				memoryNonces.Store(key, &memoryNonceEntry{expiresAt: expiresAt})
				return true, nil
			}
		}
		return false, nil
	}
	return true, nil
}