Efficient Invalid JWT Mark & Detect using expiring Bloom Filters

benchmark_test.go

package main

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// Define a custom claims structure that includes standard claims
type MyClaims struct {
	jwt.RegisteredClaims
	UserID   string `json:"user_id"`
	UserRole string `json:"role"`
}

// Global variable to hold the token string for verification benchmarks
var tokenString string

// Secret key for HMAC signing (must be a strong, cryptographically random key in production)
const hmacSecret = "super-secret-key-for-hmac-sha256-verification-32-bytes"
const hmacSecret2 = hmacSecret + "-2"

func init() {
	claims := MyClaims{
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour * 24)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			Subject:   "test-user",
		},
		UserID:   "user-12345",
		UserRole: "admin",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	var err error
	tokenString, err = token.SignedString([]byte(hmacSecret2))
	if err != nil {
		panic(fmt.Sprintf("Failed to sign token in init: %v", err))
	}
}

// Key function for verification that returns the secret
func keyFunc(token *jwt.Token) (interface{}, error) {
	// Always check the signing method to prevent algorithm switching attacks
	if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
		return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
	}
	return []byte(hmacSecret), nil
}

// Benchmark the JWT verification process using HS256 and full claims validation
func BenchmarkJWTVerification_HS256_Failed(b *testing.B) {
	var parsedToken *jwt.Token

	// Reset timer to exclude setup time
	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		token, err := jwt.ParseWithClaims(tokenString, &MyClaims{}, keyFunc)
		if err != nil {
			continue
		}
		b.Fail()
		// Store the result to prevent the compiler from optimizing the whole loop away
		parsedToken = token
	}
	// Use the result outside the loop
	_ = parsedToken
}

func BenchmarkJWTVerification_HS256_BloomFilter(b *testing.B) {
	var parsedToken *jwt.Token

	b1, cancel := NewTTLBloomFilterWithContext(context.Background(), time.Minute)
	defer cancel()

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		if b1.Test([]byte(tokenString)) {
			continue
		}
		token, err := jwt.ParseWithClaims(tokenString, &MyClaims{}, keyFunc)
		if errors.Is(err, jwt.ErrTokenSignatureInvalid) {
			b1.Add([]byte(tokenString))
			continue
		}
		b.Fail()
		// Store the result to prevent the compiler from optimizing the whole loop away
		parsedToken = token
	}
	// Use the result outside the loop
	_ = parsedToken
}

main.go

package main

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/bits-and-blooms/bloom/v3"
)

type TTLBloomFilter struct {
	*bloom.BloomFilter
	ttl       time.Duration
	lock      sync.RWMutex
	workerCtx context.Context
}

func (f *TTLBloomFilter) Add(item []byte) {
	f.lock.Lock()
	defer f.lock.Unlock()
	f.BloomFilter.Add(item)
}

func (f *TTLBloomFilter) Test(item []byte) bool {
	f.lock.RLock()
	defer f.lock.RUnlock()
	return f.BloomFilter.Test(item)
}

func (f *TTLBloomFilter) WithContext(ctx context.Context) *TTLBloomFilter {
	f.lock.RLock()
	defer f.lock.RUnlock()
	return &TTLBloomFilter{
		BloomFilter: f.Copy(),
		ttl:         f.ttl,
		workerCtx:   ctx,
	}
}

func (f *TTLBloomFilter) ExpirationWorker() {
	ctx := f.workerCtx
	if ctx == nil {
		ctx = context.Background()
	}
	ticker := time.NewTicker(f.ttl)
	resetFunc := func() {
		f.lock.Lock()
		defer f.lock.Unlock()
		f.ClearAll()
	}
	for {
		select {
		case <-ticker.C:
			resetFunc()
		case <-ctx.Done():
			resetFunc()
			return
		}
	}
}

func NewTTLBloomFilterWithContext(ctx context.Context, ttl time.Duration) (*TTLBloomFilter, context.CancelFunc) {
	ctx, cancel := context.WithCancel(ctx)
	bf := &TTLBloomFilter{
		BloomFilter: bloom.NewWithEstimates(1_000_000, 0.00001),
		ttl:         ttl,
		workerCtx:   ctx,
	}
	go bf.ExpirationWorker()
	return bf, cancel
}

func NewTTLBloomFilter(ttl time.Duration) *TTLBloomFilter {
	bf := &TTLBloomFilter{
		BloomFilter: bloom.NewWithEstimates(1_000_000, 0.00001),
		ttl:         ttl,
	}
	return bf
}

func main() {
	bf, cancel := NewTTLBloomFilterWithContext(context.Background(), 5*time.Second)
	defer cancel()
	bf.Add([]byte("test1"))
	if !bf.Test([]byte("test1")) {
		panic("test1 not found in bloom filter")
	}
	if bf.Test([]byte("test2")) {
		panic("test2 found in bloom filter")
	}
	time.Sleep(5 * time.Second)
	if bf.Test([]byte("test1")) {
		panic("test1 found in bloom filter")
	}
	fmt.Println(bf.Cap() / 1024. / 1024.)
}