feat: deduplicate proxy token creation
This commit is contained in:
@@ -15,33 +15,52 @@ type proxyTokenTarget struct {
|
|||||||
expiresAt time.Time
|
expiresAt time.Time
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type proxyTokenKey struct {
|
||||||
|
targetURL string
|
||||||
|
referer string
|
||||||
|
scope string
|
||||||
|
}
|
||||||
|
|
||||||
type proxyTokenStore struct {
|
type proxyTokenStore struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
tokens map[string]proxyTokenTarget
|
tokens map[string]proxyTokenTarget
|
||||||
|
byTarget map[proxyTokenKey]string
|
||||||
}
|
}
|
||||||
|
|
||||||
func newProxyTokenStore() *proxyTokenStore {
|
func newProxyTokenStore() *proxyTokenStore {
|
||||||
return &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) {
|
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)
|
tokenBytes := make([]byte, 32)
|
||||||
if _, err := rand.Read(tokenBytes); err != nil {
|
if _, err := rand.Read(tokenBytes); err != nil {
|
||||||
return "", fmt.Errorf("generate proxy token: %w", err)
|
return "", fmt.Errorf("generate proxy token: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
token := base64.RawURLEncoding.EncodeToString(tokenBytes)
|
token := base64.RawURLEncoding.EncodeToString(tokenBytes)
|
||||||
s.mu.Lock()
|
|
||||||
defer s.mu.Unlock()
|
|
||||||
s.pruneExpiredLocked(now)
|
|
||||||
s.tokens[token] = proxyTokenTarget{
|
s.tokens[token] = proxyTokenTarget{
|
||||||
targetURL: targetURL,
|
targetURL: targetURL,
|
||||||
referer: referer,
|
referer: referer,
|
||||||
scope: scope,
|
scope: scope,
|
||||||
expiresAt: now.Add(ttl),
|
expiresAt: now.Add(ttl),
|
||||||
}
|
}
|
||||||
|
s.byTarget[key] = token
|
||||||
return token, nil
|
return token, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -55,6 +74,7 @@ func (s *proxyTokenStore) resolve(token string, now time.Time) (proxyTokenTarget
|
|||||||
}
|
}
|
||||||
if !target.expiresAt.After(now) {
|
if !target.expiresAt.After(now) {
|
||||||
delete(s.tokens, token)
|
delete(s.tokens, token)
|
||||||
|
delete(s.byTarget, target.key())
|
||||||
return proxyTokenTarget{}, fmt.Errorf("proxy token expired")
|
return proxyTokenTarget{}, fmt.Errorf("proxy token expired")
|
||||||
}
|
}
|
||||||
return target, nil
|
return target, nil
|
||||||
@@ -64,6 +84,11 @@ func (s *proxyTokenStore) pruneExpiredLocked(now time.Time) {
|
|||||||
for token, target := range s.tokens {
|
for token, target := range s.tokens {
|
||||||
if !target.expiresAt.After(now) {
|
if !target.expiresAt.After(now) {
|
||||||
delete(s.tokens, token)
|
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}
|
||||||
|
}
|
||||||
|
|||||||
68
internal/playback/proxy_tokens_test.go
Normal file
68
internal/playback/proxy_tokens_test.go
Normal file
@@ -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))
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user