ietf-wimse-ect/refimpl/go-lang/ect/jti_cache.go

package ect

import (
	"sync"
	"time"
)

// JTICache provides replay protection by remembering seen JTIs. Safe for concurrent use.
// Use as VerifyOptions.JTISeen: cache.Seen(jti) and call cache.Add(jti) after successful verify.
type JTICache interface {
	// Seen returns true if jti was already added (replay).
	Seen(jti string) bool
	// Add records jti. Call after successful verification before accepting the token.
	Add(jti string)
}

type jtiEntry struct {
	expiresAt time.Time
}

// NewJTICache returns an in-memory JTI cache with optional max size and TTL.
// maxSize 0 means unbounded; entries are evicted after ttl.
func NewJTICache(maxSize int, ttl time.Duration) JTICache {
	return &jtiCache{
		byJti:   make(map[string]jtiEntry),
		maxSize: maxSize,
		ttl:     ttl,
	}
}

type jtiCache struct {
	mu      sync.RWMutex
	byJti   map[string]jtiEntry
	maxSize int
	ttl     time.Duration
}

func (c *jtiCache) Seen(jti string) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	e, ok := c.byJti[jti]
	if !ok {
		return false
	}
	if time.Now().After(e.expiresAt) {
		return false
	}
	return true
}

func (c *jtiCache) Add(jti string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	now := time.Now()
	// Evict expired
	for k, v := range c.byJti {
		if now.After(v.expiresAt) {
			delete(c.byJti, k)
		}
	}
	if c.maxSize > 0 && len(c.byJti) >= c.maxSize && c.byJti[jti].expiresAt.IsZero() {
		// Evict one oldest (simple: remove first we see)
		for k := range c.byJti {
			delete(c.byJti, k)
			break
		}
	}
	c.byJti[jti] = jtiEntry{expiresAt: now.Add(c.ttl)}
}