feat: add HLS manifest cache with TTL-based expiry

internal/playback/handler/manifest_cache.go

@@ -0,0 +1,109 @@
+package handler
+
+import (
+	"container/list"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+)
+
+const (
+	manifestCacheVODTTL      = 45 * time.Second
+	manifestCacheLiveTTL     = 3 * time.Second
+	manifestCacheMaxEntries  = 256
+	manifestCacheMaxBodySize = 2 << 20
+)
+
+type manifestCacheEntry struct {
+	key       string
+	body      []byte
+	headers   http.Header
+	status    int
+	expiresAt time.Time
+}
+
+type manifestCache struct {
+	mu         sync.Mutex
+	maxEntries int
+	entries    map[string]*list.Element
+	lru        *list.List
+}
+
+func newManifestCache(maxEntries int) *manifestCache {
+	if maxEntries <= 0 {
+		maxEntries = manifestCacheMaxEntries
+	}
+	return &manifestCache{
+		maxEntries: maxEntries,
+		entries:    make(map[string]*list.Element, maxEntries),
+		lru:        list.New(),
+	}
+}
+
+func manifestCacheKey(targetURL string, referer string) string {
+	return targetURL + "\x00" + referer
+}
+
+func (c *manifestCache) get(key string, now time.Time) (manifestCacheEntry, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	el := c.entries[key]
+	if el == nil {
+		return manifestCacheEntry{}, false
+	}
+	entry, ok := el.Value.(manifestCacheEntry)
+	if !ok || !now.Before(entry.expiresAt) {
+		c.remove(el)
+		return manifestCacheEntry{}, false
+	}
+	c.lru.MoveToFront(el)
+	entry.body = append([]byte(nil), entry.body...)
+	entry.headers = entry.headers.Clone()
+	return entry, true
+}
+
+func (c *manifestCache) set(key string, status int, headers http.Header, body []byte, now time.Time) bool {
+	if status < http.StatusOK || status >= http.StatusMultipleChoices || len(body) > manifestCacheMaxBodySize {
+		return false
+	}
+
+	ttl := manifestCacheVODTTL
+	bodyText := string(body)
+	if strings.Contains(bodyText, "#EXTINF") && !strings.Contains(bodyText, "#EXT-X-ENDLIST") {
+		ttl = manifestCacheLiveTTL
+	}
+	entry := manifestCacheEntry{
+		key:       key,
+		body:      append([]byte(nil), body...),
+		headers:   headers.Clone(),
+		status:    status,
+		expiresAt: now.Add(ttl),
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if el := c.entries[key]; el != nil {
+		el.Value = entry
+		c.lru.MoveToFront(el)
+		return true
+	}
+	c.entries[key] = c.lru.PushFront(entry)
+	for len(c.entries) > c.maxEntries {
+		back := c.lru.Back()
+		if back == nil {
+			break
+		}
+		c.remove(back)
+	}
+	return true
+}
+
+func (c *manifestCache) remove(el *list.Element) {
+	entry, ok := el.Value.(manifestCacheEntry)
+	if ok {
+		delete(c.entries, entry.key)
+	}
+	c.lru.Remove(el)
+}

internal/playback/handler/manifest_cache_test.go

@@ -0,0 +1,122 @@
+package handler
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+type manifestRoundTripper struct {
+	calls int
+	body  string
+}
+
+func (rt *manifestRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	rt.calls++
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Header: http.Header{
+			"Content-Type":   {"application/vnd.apple.mpegurl"},
+			"Content-Length": {"48"},
+		},
+		Body: io.NopCloser(strings.NewReader(rt.body)),
+	}, nil
+}
+
+func TestManifestCacheTTLSizeAndLRU(t *testing.T) {
+	now := time.Unix(100, 0)
+	cache := newManifestCache(2)
+	headers := http.Header{"Content-Type": {"application/vnd.apple.mpegurl"}}
+	vod := []byte("#EXTM3U\n#EXTINF:4,\nsegment.ts\n#EXT-X-ENDLIST\n")
+	live := []byte("#EXTM3U\n#EXTINF:4,\nsegment.ts\n")
+
+	if !cache.set("vod", http.StatusOK, headers, vod, now) {
+		t.Fatal("VOD manifest was not cached")
+	}
+	if !cache.set("live", http.StatusOK, headers, live, now) {
+		t.Fatal("live manifest was not cached")
+	}
+	cache.get("vod", now)
+	cache.set("new", http.StatusOK, headers, vod, now)
+
+	if _, ok := cache.get("live", now); ok {
+		t.Fatal("least recently used manifest was not evicted")
+	}
+	if _, ok := cache.get("vod", now.Add(manifestCacheLiveTTL+time.Second)); !ok {
+		t.Fatal("VOD manifest expired at the live TTL")
+	}
+	if cache.set("large", http.StatusOK, headers, make([]byte, manifestCacheMaxBodySize+1), now) {
+		t.Fatal("oversized manifest was cached")
+	}
+}
+
+func TestReadBoundedPlaylistRejectsOversizedBody(t *testing.T) {
+	_, err := readBoundedPlaylist(bytes.NewReader(make([]byte, manifestCacheMaxBodySize+1)))
+	if err == nil {
+		t.Fatal("readBoundedPlaylist accepted an oversized body")
+	}
+}
+
+func TestCachedRawPlaylistIsRewrittenWithResponseTokens(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	svc := &rewritePlaybackService{}
+	h := &PlaybackHandler{svc: svc}
+	body := []byte("#EXTM3U\n#EXTINF:4,\nsegment.ts\n#EXT-X-ENDLIST\n")
+	headers := http.Header{"Content-Type": {"application/vnd.apple.mpegurl"}, "Content-Length": {"48"}}
+
+	responses := make([]string, 0, 2)
+	for range 2 {
+		recorder := httptest.NewRecorder()
+		ctx, _ := gin.CreateTestContext(recorder)
+		h.writeProxyPlaylist(ctx, http.StatusOK, headers, body, "https://cdn.example.test/master.m3u8", "")
+		responses = append(responses, recorder.Body.String())
+		if recorder.Header().Get("Content-Length") == "48" {
+			t.Fatal("rewritten playlist retained upstream Content-Length")
+		}
+	}
+
+	if responses[0] == responses[1] {
+		t.Fatalf("rewritten responses reused tokenized output: %q", responses[0])
+	}
+	if !strings.Contains(responses[0], "token=token-1") || !strings.Contains(responses[1], "token=token-2") {
+		t.Fatalf("responses did not contain separately issued tokens: %#v", responses)
+	}
+}
+
+func TestHandleProxyStreamServesCachedRawManifest(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	svc := &rewritePlaybackService{targetURL: "https://203.0.113.10/master.m3u8"}
+	rt := &manifestRoundTripper{body: "#EXTM3U\n#EXTINF:4,\nsegment.ts\n#EXT-X-ENDLIST\n"}
+	h := &PlaybackHandler{
+		svc:             svc,
+		streamingClient: &http.Client{Transport: rt},
+		manifestCache:   newManifestCache(2),
+	}
+	router := gin.New()
+	router.GET("/watch/proxy/stream", h.HandleProxyStream)
+
+	responses := make([]string, 0, 2)
+	for range 2 {
+		recorder := httptest.NewRecorder()
+		request := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/watch/proxy/stream?token=opaque", nil)
+		router.ServeHTTP(recorder, request)
+		if recorder.Code != http.StatusOK {
+			t.Fatalf("status = %d, want %d", recorder.Code, http.StatusOK)
+		}
+		responses = append(responses, recorder.Body.String())
+	}
+
+	if rt.calls != 1 {
+		t.Fatalf("upstream calls = %d, want 1", rt.calls)
+	}
+	if responses[0] == responses[1] {
+		t.Fatalf("cached raw manifest reused rewritten tokens: %#v", responses)
+	}
+}