sec: add jti and replay-protect subscribe tokens

internal/auth/token.go

@@ -2,33 +2,59 @@ package auth
 
 import (
 	"crypto/hmac"
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
 var (
 	ErrTokenMalformed = errors.New("token malformed")
 	ErrTokenExpired   = errors.New("token expired")
 	ErrTokenInvalid   = errors.New("token invalid")
+	ErrTokenReplayed  = errors.New("token already used")
 )
 
-// SignToken returns "<expMillis>:<base64mac>".
+// Token format: "<expMs>:<jti>:<base64mac>"
+// jti is a random 16-byte hex string (32 chars) embedded in the HMAC payload
+// so it can't be swapped out without breaking the signature. ReplayCache
+// records jti -> expiry; a second VerifyTokenAgainst call with the same
+// jti returns ErrTokenReplayed.
+
+// SignToken signs a one-time-use token. Each call generates a fresh jti, so
+// distinct calls with the same args still produce distinct tokens.
 func SignToken(secret, socketID, channel string, expiry time.Time) (string, error) {
 	expMs := expiry.UnixMilli()
-	payload := fmt.Sprintf("%d|%s|%s", expMs, socketID, channel)
+	jtiBytes := make([]byte, 16)
+	if _, err := rand.Read(jtiBytes); err != nil {
+		return "", err
+	}
+	jti := hex.EncodeToString(jtiBytes)
+	payload := fmt.Sprintf("%d|%s|%s|%s", expMs, socketID, channel, jti)
 	mac := hmac.New(sha256.New, []byte(secret))
 	mac.Write([]byte(payload))
-	return strconv.FormatInt(expMs, 10) + ":" + base64.RawURLEncoding.EncodeToString(mac.Sum(nil)), nil
+	return strconv.FormatInt(expMs, 10) + ":" + jti + ":" + base64.RawURLEncoding.EncodeToString(mac.Sum(nil)), nil
 }
 
+// VerifyToken validates a token without replay protection. Provided for
+// backward-compatibility; production callers should construct a ReplayCache
+// and call VerifyTokenAgainst.
 func VerifyToken(secret, socketID, channel, tok string) error {
-	parts := strings.SplitN(tok, ":", 2)
-	if len(parts) != 2 {
+	return VerifyTokenAgainst(secret, socketID, channel, tok, nil)
+}
+
+// VerifyTokenAgainst is VerifyToken with optional replay protection.
+// When cache is non-nil and the token verifies cleanly, the jti is recorded;
+// a subsequent call with the same jti returns ErrTokenReplayed.
+func VerifyTokenAgainst(secret, socketID, channel, tok string, cache *ReplayCache) error {
+	parts := strings.SplitN(tok, ":", 3)
+	if len(parts) != 3 {
 		return ErrTokenMalformed
 	}
 	expMs, err := strconv.ParseInt(parts[0], 10, 64)
@@ -38,14 +64,69 @@ func VerifyToken(secret, socketID, channel, tok string) error {
 	if time.Now().UnixMilli() > expMs {
 		return ErrTokenExpired
 	}
-	sig, err := base64.RawURLEncoding.DecodeString(parts[1])
+	jti := parts[1]
+	if jti == "" {
+		return ErrTokenMalformed
+	}
+	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
 	if err != nil {
 		return ErrTokenMalformed
 	}
 	mac := hmac.New(sha256.New, []byte(secret))
-	_, _ = fmt.Fprintf(mac, "%d|%s|%s", expMs, socketID, channel)
+	_, _ = fmt.Fprintf(mac, "%d|%s|%s|%s", expMs, socketID, channel, jti)
 	if !hmac.Equal(sig, mac.Sum(nil)) {
 		return ErrTokenInvalid
 	}
+	if cache != nil {
+		if !cache.CheckAndRecord(jti, time.UnixMilli(expMs)) {
+			return ErrTokenReplayed
+		}
+	}
 	return nil
 }
+
+// ReplayCache records token jti -> expiry. Safe for concurrent use. Memory
+// stays bounded by Sweep, which removes expired entries; callers that don't
+// run Sweep periodically will accumulate memory at the rate of issued
+// tokens until the next Sweep.
+type ReplayCache struct {
+	mu   sync.Mutex
+	seen map[string]time.Time
+}
+
+func NewReplayCache() *ReplayCache {
+	return &ReplayCache{seen: map[string]time.Time{}}
+}
+
+// CheckAndRecord returns true the first time it sees jti, false thereafter.
+func (c *ReplayCache) CheckAndRecord(jti string, exp time.Time) bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if _, ok := c.seen[jti]; ok {
+		return false
+	}
+	c.seen[jti] = exp
+	return true
+}
+
+// Sweep removes entries whose expiry has passed. Returns the count removed.
+func (c *ReplayCache) Sweep() int {
+	now := time.Now()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	swept := 0
+	for jti, exp := range c.seen {
+		if !now.Before(exp) {
+			delete(c.seen, jti)
+			swept++
+		}
+	}
+	return swept
+}
+
+// Len reports the current number of cached entries (for tests / metrics).
+func (c *ReplayCache) Len() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.seen)
+}

internal/auth/token_test.go

@@ -41,3 +41,47 @@ func TestVerifyTokenTampered(t *testing.T) {
 		t.Fatal("expected tamper error")
 	}
 }
+
+func TestVerifyTokenAgainstCachePreventsReplay(t *testing.T) {
+	secret := "s"
+	tok, _ := SignToken(secret, "sock1", "private-x", time.Now().Add(time.Minute))
+	cache := NewReplayCache()
+	if err := VerifyTokenAgainst(secret, "sock1", "private-x", tok, cache); err != nil {
+		t.Fatalf("first verify: %v", err)
+	}
+	if err := VerifyTokenAgainst(secret, "sock1", "private-x", tok, cache); err == nil {
+		t.Fatal("expected ErrTokenReplayed on second use, got nil")
+	}
+}
+
+func TestVerifyTokenAgainstNilCache(t *testing.T) {
+	secret := "s"
+	tok, _ := SignToken(secret, "sock1", "private-x", time.Now().Add(time.Minute))
+	if err := VerifyTokenAgainst(secret, "sock1", "private-x", tok, nil); err != nil {
+		t.Fatalf("first verify with nil cache: %v", err)
+	}
+	if err := VerifyTokenAgainst(secret, "sock1", "private-x", tok, nil); err != nil {
+		t.Fatalf("second verify with nil cache: %v", err)
+	}
+}
+
+func TestReplayCacheSweepRemovesExpired(t *testing.T) {
+	c := NewReplayCache()
+	c.CheckAndRecord("expired", time.Now().Add(-time.Minute))
+	c.CheckAndRecord("alive", time.Now().Add(time.Minute))
+	if got := c.Sweep(); got != 1 {
+		t.Errorf("Sweep removed %d, want 1", got)
+	}
+	if got, want := c.Len(), 1; got != want {
+		t.Errorf("post-Sweep Len = %d, want %d", got, want)
+	}
+}
+
+func TestSignTokenJtiUnique(t *testing.T) {
+	secret := "s"
+	a, _ := SignToken(secret, "sock1", "x", time.Now().Add(time.Minute))
+	b, _ := SignToken(secret, "sock1", "x", time.Now().Add(time.Minute))
+	if a == b {
+		t.Fatal("two SignToken calls with identical args produced identical tokens — jti collision or absent")
+	}
+}

internal/conn/conn.go

@@ -9,6 +9,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/EthanY33/wirefan/internal/auth"
 	"github.com/EthanY33/wirefan/internal/fanout"
 	"github.com/EthanY33/wirefan/internal/hub"
 	"github.com/EthanY33/wirefan/internal/metrics"
@@ -37,6 +38,7 @@ type Conn struct {
 	send          chan []byte
 	registry      registry.Registry
 	signingSecret string
+	replayCache   *auth.ReplayCache
 	fanout        fanout.Fanout
 	rateLimit     *ratelimit.Limiter // per-API-key bucket; shared across all conns owned by the key
 	connRate      *rate.Limiter      // per-conn bucket; bounds a single socket's throughput
@@ -70,14 +72,19 @@ func (c *Conn) CloseFrame(code websocket.StatusCode, reason string) {
 }
 
 // Run owns the conn for its lifetime. Returns when ctx is canceled or peer disconnects.
-func Run(ctx context.Context, ws *websocket.Conn, socketID, apiKeyID string, reg registry.Registry, signingSecret string, fan fanout.Fanout, rl *ratelimit.Limiter, pol Policy, h *hub.Hub) error {
+//
+// replayCache may be nil; when nil, subscribe-token replay protection is
+// disabled (used by tests). Production callers pass a process-wide cache so
+// a leaked subscribe token cannot be reused within its 5-minute window.
+func Run(ctx context.Context, ws *websocket.Conn, socketID, apiKeyID string, reg registry.Registry, signingSecret string, replayCache *auth.ReplayCache, fan fanout.Fanout, rl *ratelimit.Limiter, pol Policy, h *hub.Hub) error {
 	c := &Conn{
 		ws:            ws,
 		socketID:      socketID,
 		apiKeyID:      apiKeyID,
 		send:          make(chan []byte, sendChanSize),
 		registry:      reg,
 		signingSecret: signingSecret,
+		replayCache:   replayCache,
 		fanout:        fan,
 		rateLimit:     rl,
 		connRate:      rate.NewLimiter(rate.Limit(defaultConnPublishRate), defaultConnPublishBurst),

internal/conn/conn_test.go

@@ -27,7 +27,7 @@ func TestConnectedMessageSent(t *testing.T) {
 	handler := func(c *websocket.Conn) {
 		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 		defer cancel()
-		_ = Run(ctx, c, "01HTEST", "test-key", registry.NewSyncMap(), "test-signing-secret", fanout.NewPerConn(), rl, PolicyDisconnect{}, hub.New())
+		_ = Run(ctx, c, "01HTEST", "test-key", registry.NewSyncMap(), "test-signing-secret", nil, fanout.NewPerConn(), rl, PolicyDisconnect{}, hub.New())
 	}
 
 	srv := httptest.NewServer(websocketHandler(handler))

internal/conn/handler.go

@@ -81,8 +81,12 @@ func (c *Conn) handleSubscribe(msg incoming) {
 		return
 	}
 	if strings.HasPrefix(msg.Channel, "private-") {
-		if err := auth.VerifyToken(c.signingSecret, c.socketID, msg.Channel, msg.Token); err != nil {
+		if err := auth.VerifyTokenAgainst(c.signingSecret, c.socketID, msg.Channel, msg.Token, c.replayCache); err != nil {
 			metrics.AuthFails.Inc()
+			if errors.Is(err, auth.ErrTokenReplayed) {
+				c.sendError("AUTH_REPLAYED", "token already used")
+				return
+			}
 			c.sendError("AUTH_FAILED", "invalid token")
 			return
 		}

internal/conn/handler_test.go

@@ -31,7 +31,7 @@ func newTestConn(t *testing.T, signingSecret string) (*websocket.Conn, string) {
 		}
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
-		_ = Run(ctx, c, socketID, "test-key", registry.NewSyncMap(), signingSecret, fanout.NewPerConn(), rl, PolicyDisconnect{}, hub.New())
+		_ = Run(ctx, c, socketID, "test-key", registry.NewSyncMap(), signingSecret, nil, fanout.NewPerConn(), rl, PolicyDisconnect{}, hub.New())
 	})
 	srv := httptest.NewServer(handler)
 	t.Cleanup(srv.Close)

internal/server/leak_test.go

@@ -33,6 +33,7 @@ func TestNoGoroutineLeakAfterChurn(t *testing.T) {
 		[]string{"*"},
 		registry.NewSyncMap(),
 		"test-signing-secret",
+		nil,
 		fanout.NewPerConn(),
 		rl,
 		conn.PolicyDisconnect{},

internal/server/server.go

@@ -8,6 +8,7 @@ import (
 	"net/http/pprof"
 	"time"
 
+	"github.com/EthanY33/wirefan/internal/auth"
 	"github.com/EthanY33/wirefan/internal/conn"
 	"github.com/EthanY33/wirefan/internal/fanout"
 	"github.com/EthanY33/wirefan/internal/hub"
@@ -33,34 +34,37 @@ type Config struct {
 }
 
 type Server struct {
-	cfg      Config
-	health   *HealthHandler
-	mux      *http.ServeMux
-	adminMux *http.ServeMux
-	srv      *http.Server
-	adminSrv *http.Server
-	store    store.Store
-	hub      *hub.Hub
+	cfg         Config
+	health      *HealthHandler
+	mux         *http.ServeMux
+	adminMux    *http.ServeMux
+	srv         *http.Server
+	adminSrv    *http.Server
+	store       store.Store
+	hub         *hub.Hub
+	replayCache *auth.ReplayCache
 }
 
 // New builds the public and admin muxes. The admin listener is created
 // only when cfg.AdminAddr is non-empty.
 func New(cfg Config, st store.Store, adminToken string, reg registry.Registry, signingSecret string, fan fanout.Fanout, rl *ratelimit.Limiter, pol conn.Policy, h *hub.Hub) *Server {
+	rc := auth.NewReplayCache()
 	s := &Server{
-		cfg:      cfg,
-		health:   NewHealthHandler(),
-		mux:      http.NewServeMux(),
-		adminMux: http.NewServeMux(),
-		store:    st,
-		hub:      h,
+		cfg:         cfg,
+		health:      NewHealthHandler(),
+		mux:         http.NewServeMux(),
+		adminMux:    http.NewServeMux(),
+		store:       st,
+		hub:         h,
+		replayCache: rc,
 	}
 
 	rest := NewRestHandler(st, adminToken, signingSecret)
 
 	// Public listener: health, /v1/connect (WS), /v1/auth/sign, static client.
 	s.mux.Handle("/v1/health", s.health)
 	rest.RegisterPublic(s.mux)
-	s.mux.Handle("/v1/connect", NewUpgradeHandler(st, cfg.AllowedOrigins, reg, signingSecret, fan, rl, pol, h))
+	s.mux.Handle("/v1/connect", NewUpgradeHandler(st, cfg.AllowedOrigins, reg, signingSecret, rc, fan, rl, pol, h))
 	s.mux.Handle("/", http.FileServerFS(web.Files))
 
 	// Admin listener: metrics, pprof, key management. All gated by
@@ -81,6 +85,11 @@ func New(cfg Config, st store.Store, adminToken string, reg registry.Registry, s
 	return s
 }
 
+// ReplayCache exposes the per-Server token replay cache so the caller can
+// run a periodic Sweep goroutine. Public so cmd/wirefan/main.go can drive
+// the sweeper without exporting a separate accessor.
+func (s *Server) ReplayCache() *auth.ReplayCache { return s.replayCache }
+
 func (s *Server) Run(ctx context.Context) error {
 	errc := make(chan error, 2)
 	go func() {
@@ -97,6 +106,7 @@ func (s *Server) Run(ctx context.Context) error {
 			}
 		}()
 	}
+	go s.sweepReplayCache(ctx)
 
 	select {
 	case err := <-errc:
@@ -113,3 +123,20 @@ func (s *Server) Run(ctx context.Context) error {
 	}
 	return s.srv.Shutdown(shutdownCtx)
 }
+
+// sweepReplayCache evicts expired token jti entries every minute. Memory in
+// the cache is bounded by the issuance rate * token lifetime (5 minutes by
+// default), so a sweep cadence of one minute gives at most ~5 minutes of
+// expired entries before reclamation. Loops until ctx is canceled.
+func (s *Server) sweepReplayCache(ctx context.Context) {
+	t := time.NewTicker(time.Minute)
+	defer t.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-t.C:
+			s.replayCache.Sweep()
+		}
+	}
+}

internal/server/shutdown_test.go

@@ -28,7 +28,7 @@ func TestDrainClosesAllConnections(t *testing.T) {
 
 	upgrader := NewUpgradeHandler(
 		s, []string{"*"}, registry.NewSyncMap(), "test-secret",
-		fanout.NewPerConn(), rl, conn.PolicyDisconnect{}, h,
+		nil, fanout.NewPerConn(), rl, conn.PolicyDisconnect{}, h,
 	)
 	srv := httptest.NewServer(upgrader)
 	defer srv.Close()