From 622418f96c02187af26ad68099bf84f9d6c28a5f Mon Sep 17 00:00:00 2001
From: mkelvers <mikkelelvers@outlook.com>
Date: Tue, 16 Jun 2026 01:04:21 +0200
Subject: [PATCH] feat: deduplicate proxy token creation

---
 internal/playback/proxy_tokens.go      | 37 +++++++++++---
 internal/playback/proxy_tokens_test.go | 68 ++++++++++++++++++++++++++
 2 files changed, 99 insertions(+), 6 deletions(-)
 create mode 100644 internal/playback/proxy_tokens_test.go

diff --git a/internal/playback/proxy_tokens.go b/internal/playback/proxy_tokens.go
index 01365cf..f940ca4 100644
--- a/internal/playback/proxy_tokens.go
+++ b/internal/playback/proxy_tokens.go
@@ -15,33 +15,52 @@ type proxyTokenTarget struct {
 	expiresAt time.Time
 }
 
+type proxyTokenKey struct {
+	targetURL string
+	referer   string
+	scope     string
+}
+
 type proxyTokenStore struct {
-	mu     sync.Mutex
-	tokens map[string]proxyTokenTarget
+	mu       sync.Mutex
+	tokens   map[string]proxyTokenTarget
+	byTarget map[proxyTokenKey]string
 }
 
 func newProxyTokenStore() *proxyTokenStore {
 	return &proxyTokenStore{
-		tokens: make(map[string]proxyTokenTarget),
+		tokens:   make(map[string]proxyTokenTarget),
+		byTarget: make(map[proxyTokenKey]string),
 	}
 }
 
 func (s *proxyTokenStore) create(targetURL, referer, scope string, ttl time.Duration, now time.Time) (string, error) {
+	key := proxyTokenKey{targetURL: targetURL, referer: referer, scope: scope}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.pruneExpiredLocked(now)
+	if token, ok := s.byTarget[key]; ok {
+		if target, ok := s.tokens[token]; ok && target.expiresAt.After(now) {
+			return token, nil
+		}
+		delete(s.byTarget, key)
+	}
+
 	tokenBytes := make([]byte, 32)
 	if _, err := rand.Read(tokenBytes); err != nil {
 		return "", fmt.Errorf("generate proxy token: %w", err)
 	}
 
 	token := base64.RawURLEncoding.EncodeToString(tokenBytes)
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.pruneExpiredLocked(now)
 	s.tokens[token] = proxyTokenTarget{
 		targetURL: targetURL,
 		referer:   referer,
 		scope:     scope,
 		expiresAt: now.Add(ttl),
 	}
+	s.byTarget[key] = token
 	return token, nil
 }
 
@@ -55,6 +74,7 @@ func (s *proxyTokenStore) resolve(token string, now time.Time) (proxyTokenTarget
 	}
 	if !target.expiresAt.After(now) {
 		delete(s.tokens, token)
+		delete(s.byTarget, target.key())
 		return proxyTokenTarget{}, fmt.Errorf("proxy token expired")
 	}
 	return target, nil
@@ -64,6 +84,11 @@ func (s *proxyTokenStore) pruneExpiredLocked(now time.Time) {
 	for token, target := range s.tokens {
 		if !target.expiresAt.After(now) {
 			delete(s.tokens, token)
+			delete(s.byTarget, target.key())
 		}
 	}
 }
+
+func (t proxyTokenTarget) key() proxyTokenKey {
+	return proxyTokenKey{targetURL: t.targetURL, referer: t.referer, scope: t.scope}
+}
diff --git a/internal/playback/proxy_tokens_test.go b/internal/playback/proxy_tokens_test.go
new file mode 100644
index 0000000..f62722b
--- /dev/null
+++ b/internal/playback/proxy_tokens_test.go
@@ -0,0 +1,68 @@
+package playback
+
+import (
+	"testing"
+	"time"
+)
+
+func TestProxyTokenStoreReusesUnexpiredTokenForSameTarget(t *testing.T) {
+	store := newProxyTokenStore()
+	now := time.Date(2026, 6, 16, 12, 0, 0, 0, time.UTC)
+
+	first, err := store.create("https://cdn.example.test/seg.ts", "https://referer.example.test", "stream", time.Hour, now)
+	if err != nil {
+		t.Fatalf("create first token: %v", err)
+	}
+	second, err := store.create("https://cdn.example.test/seg.ts", "https://referer.example.test", "stream", time.Hour, now.Add(time.Minute))
+	if err != nil {
+		t.Fatalf("create second token: %v", err)
+	}
+
+	if first != second {
+		t.Fatalf("tokens differ for same target: %q != %q", first, second)
+	}
+	if len(store.tokens) != 1 {
+		t.Fatalf("token count = %d, want 1", len(store.tokens))
+	}
+}
+
+func TestProxyTokenStoreCreatesNewTokenAfterExpiry(t *testing.T) {
+	store := newProxyTokenStore()
+	now := time.Date(2026, 6, 16, 12, 0, 0, 0, time.UTC)
+
+	first, err := store.create("https://cdn.example.test/seg.ts", "https://referer.example.test", "stream", time.Hour, now)
+	if err != nil {
+		t.Fatalf("create first token: %v", err)
+	}
+	second, err := store.create("https://cdn.example.test/seg.ts", "https://referer.example.test", "stream", time.Hour, now.Add(time.Hour))
+	if err != nil {
+		t.Fatalf("create second token: %v", err)
+	}
+
+	if first == second {
+		t.Fatalf("token was reused after expiry: %q", first)
+	}
+	if len(store.tokens) != 1 {
+		t.Fatalf("token count = %d, want 1", len(store.tokens))
+	}
+}
+
+func TestProxyTokenStoreCleansReverseIndexOnResolveExpiry(t *testing.T) {
+	store := newProxyTokenStore()
+	now := time.Date(2026, 6, 16, 12, 0, 0, 0, time.UTC)
+
+	token, err := store.create("https://cdn.example.test/seg.ts", "https://referer.example.test", "stream", time.Hour, now)
+	if err != nil {
+		t.Fatalf("create token: %v", err)
+	}
+	if _, err := store.resolve(token, now.Add(time.Hour)); err == nil {
+		t.Fatal("resolve expired token unexpectedly succeeded")
+	}
+
+	if len(store.tokens) != 0 {
+		t.Fatalf("token count = %d, want 0", len(store.tokens))
+	}
+	if len(store.byTarget) != 0 {
+		t.Fatalf("target index count = %d, want 0", len(store.byTarget))
+	}
+}