diff --git a/internal/api/handlers/downloads.go b/internal/api/handlers/downloads.go index d65bb441..546591cb 100644 --- a/internal/api/handlers/downloads.go +++ b/internal/api/handlers/downloads.go @@ -81,11 +81,11 @@ func (h *DownloadHandler) SetTransferRegistry(registry *transfers.Registry) { // Sessionless (downloads have no stream session), so only user-kind kills apply; // the request entry time predates any future revocation, which is exactly the // user-kill cutoff contract. Returns a stop func the caller must defer. -func (h *DownloadHandler) cutDownloadOnUserRevocation(w http.ResponseWriter, userID int) func() { +func (h *DownloadHandler) cutDownloadOnUserRevocation(ctx context.Context, w http.ResponseWriter, userID int, startedAt time.Time) func() { if h.revocation == nil { return func() {} } - return h.revocation.WatchAndCut(w, "", userID, time.Now()) + return h.revocation.WatchAndCutContext(ctx, w, "", userID, startedAt) } // downloadRequest represents the JSON body for POST /downloads. @@ -441,15 +441,17 @@ func (h *DownloadHandler) HandleDownloadFile(w http.ResponseWriter, r *http.Requ profileID, deviceID, deviceName, _ := managedIdentity(r) filter := requestAccessFilter(r) + startedAt := time.Now() + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) // Full media downloads outlive the server's absolute WriteTimeout; roll // the write deadline with progress instead. - sw := httpstream.NewRollingDeadlineWriter(w) + sw := httpstream.NewRollingDeadlineWriterCtx(r.Context(), w) transfer := h.newTransfer(r, userID, profileID, deviceName, "native_download") defer h.transfers.End(transfer.ID) metered := playback.NewSessionMeteredWriter(sw, h.transfers, transfer.ID) defer func() { _ = metered.Close() }() // In-flight kill switch: a user revocation cuts this pour mid-transfer. - stop := h.cutDownloadOnUserRevocation(metered, userID) + stop := h.cutDownloadOnUserRevocation(r.Context(), metered, userID, startedAt) defer stop() if err := h.svc.ServeFile(r.Context(), metered, r, userID, profileID, deviceID, id, filter, transfer); err != nil { @@ -497,12 +499,14 @@ func (h *DownloadHandler) HandleDirectDownload(w http.ResponseWriter, r *http.Re filter := requestAccessFilter(r) profileID, _, deviceName, _ := managedIdentity(r) + startedAt := time.Now() + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) transfer := h.newTransfer(r, userID, profileID, deviceName, "native_direct") defer h.transfers.End(transfer.ID) metered := playback.NewSessionMeteredWriter(w, h.transfers, transfer.ID) defer func() { _ = metered.Close() }() // In-flight kill switch: a user revocation cuts this pour mid-transfer. - stop := h.cutDownloadOnUserRevocation(metered, userID) + stop := h.cutDownloadOnUserRevocation(r.Context(), metered, userID, startedAt) defer stop() if err := h.svc.ServeDirect(r.Context(), metered, r, userID, fileID, r.URL.Query().Get("format"), filter, transfer); err != nil { diff --git a/internal/api/handlers/ebook_convert_serve.go b/internal/api/handlers/ebook_convert_serve.go index 19bef7a8..f01f3ddc 100644 --- a/internal/api/handlers/ebook_convert_serve.go +++ b/internal/api/handlers/ebook_convert_serve.go @@ -157,7 +157,7 @@ func serveConvertedEpub(w http.ResponseWriter, r *http.Request, file *models.Med name := strings.TrimSuffix(filepath.Base(file.FilePath), filepath.Ext(file.FilePath)) + ".epub" setConvertedEpubHeaders(w, file, key) - http.ServeContent(httpstream.NewRollingDeadlineWriter(w), r, name, stat.ModTime(), f) + http.ServeContent(httpstream.NewRollingDeadlineWriterCtx(r.Context(), w), r, name, stat.ModTime(), f) return nil } diff --git a/internal/api/handlers/ebook_reader.go b/internal/api/handlers/ebook_reader.go index 87e7fe84..1cbd77fc 100644 --- a/internal/api/handlers/ebook_reader.go +++ b/internal/api/handlers/ebook_reader.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "log/slog" "mime" "net/http" "os" @@ -20,8 +21,12 @@ import ( apimw "github.com/Silo-Server/silo-server/internal/api/middleware" "github.com/Silo-Server/silo-server/internal/catalog" + "github.com/Silo-Server/silo-server/internal/clientip" "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/models" + "github.com/Silo-Server/silo-server/internal/playback" + "github.com/Silo-Server/silo-server/internal/streamrevoke" + "github.com/Silo-Server/silo-server/internal/transfers" ) // EbookReaderProgress is the persisted reader position for one profile and ebook. @@ -100,6 +105,8 @@ type EbookReaderHandler struct { // Conversion, when set and enabled, transparently converts Kindle-family // files to EPUB on read. Nil means the feature is off. Conversion *EbookConversion + Transfers *transfers.Registry + Revocation *streamrevoke.Store } func NewEbookReaderHandler(authorizer *MediaFileAuthorizer) *EbookReaderHandler { @@ -116,6 +123,10 @@ func (h *EbookReaderHandler) HandleReadFile(w http.ResponseWriter, r *http.Reque writeError(w, http.StatusServiceUnavailable, "unavailable", "Ebook reader is not configured") return } + startedAt := time.Now() + // TODO(#305, A3): use the access token's iat as the revocation cutoff + // identity once credential identity is unified across serve surfaces. + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) contentID := strings.TrimSpace(chi.URLParam(r, "content_id")) fileID, err := strconv.Atoi(chi.URLParam(r, "file_id")) @@ -134,7 +145,43 @@ func (h *EbookReaderHandler) HandleReadFile(w http.ResponseWriter, r *http.Reque return } - if err := h.serveEbook(w, r, file); err != nil { + userID := apimw.GetUserID(r.Context()) + if h.Revocation != nil && h.Revocation.Refuse(w, "", userID, startedAt) { + return + } + + route := "ebook_read" + if h.Conversion.active(r.Context()) && isKindleReaderFile(file) { + route = "ebook_convert" + } + transfer := transfers.Transfer{ + ID: uuid.NewString(), + UserID: userID, + ProfileID: apimw.GetProfileID(r.Context()), + MediaFileID: file.ID, + Route: route, + ClientIP: clientip.FromContext(r.Context()), + ClientName: r.UserAgent(), + StartedAt: startedAt, + } + if h.Transfers != nil { + if err := h.Transfers.Begin(transfer); err != nil { + slog.DebugContext(r.Context(), "ebook transfer not monitored", + "component", "ebooks", + "transfer_id", transfer.ID, + "error", err, + ) + } + } + defer h.Transfers.End(transfer.ID) + metered := playback.NewSessionMeteredWriter(w, h.Transfers, transfer.ID) + defer func() { _ = metered.Close() }() + if h.Revocation != nil { + stop := h.Revocation.WatchAndCutContext(r.Context(), metered, "", userID, startedAt) + defer stop() + } + + if err := h.serveEbook(metered, r, file); err != nil { if errors.Is(err, catalog.ErrItemNotFound) { writeError(w, http.StatusNotFound, "not_found", "Ebook file not found") return @@ -623,7 +670,7 @@ func serveEbookInline(w http.ResponseWriter, r *http.Request, file *models.Media // Ebook files are served inline on the API origin (and reachable via the // ?token= query fallback), so browsers must never content-sniff them. w.Header().Set("X-Content-Type-Options", "nosniff") - http.ServeContent(httpstream.NewRollingDeadlineWriter(w), r, stat.Name(), stat.ModTime(), f) + http.ServeContent(httpstream.NewRollingDeadlineWriterCtx(r.Context(), w), r, stat.Name(), stat.ModTime(), f) return nil } diff --git a/internal/api/handlers/ebook_reader_test.go b/internal/api/handlers/ebook_reader_test.go index 063f0bf8..9ce10887 100644 --- a/internal/api/handlers/ebook_reader_test.go +++ b/internal/api/handlers/ebook_reader_test.go @@ -4,11 +4,13 @@ import ( "context" "encoding/json" "errors" + "io" "mime" "net/http" "net/http/httptest" "os" "strings" + "sync" "testing" "time" @@ -18,6 +20,8 @@ import ( "github.com/Silo-Server/silo-server/internal/auth" "github.com/Silo-Server/silo-server/internal/catalog" "github.com/Silo-Server/silo-server/internal/models" + "github.com/Silo-Server/silo-server/internal/streamrevoke" + "github.com/Silo-Server/silo-server/internal/transfers" ) func newEbookReaderAuthRequest(method, path string) *http.Request { @@ -202,6 +206,163 @@ func TestEbookReaderServesEbookInlineWithRangeSupport(t *testing.T) { } } +type blockingEbookWriter struct { + header http.Header + reached chan struct{} + release chan struct{} + once sync.Once + written int +} + +func (w *blockingEbookWriter) Header() http.Header { return w.header } +func (w *blockingEbookWriter) WriteHeader(int) {} +func (w *blockingEbookWriter) SetWriteDeadline(time.Time) error { + return nil +} +func (w *blockingEbookWriter) Write(p []byte) (int, error) { + w.written += len(p) + if w.written >= (1<<20)+(64<<10) { + w.once.Do(func() { close(w.reached) }) + <-w.release + } + return len(p), nil +} + +func TestEbookReaderTracksBytesWhilePourIsActiveAndFlushesBeforeRemoval(t *testing.T) { + filePath := writePlaybackTestMediaFile(t, "large.epub") + if err := os.WriteFile(filePath, make([]byte, 2<<20), 0o644); err != nil { + t.Fatal(err) + } + registry := transfers.New() + handler := NewEbookReaderHandler(&MediaFileAuthorizer{ + FileResolver: stubMediaFileResolver{file: &models.MediaFile{ + ID: 42, ContentID: "ebook-1", FilePath: filePath, Container: "epub", BaseType: "ebook", + }}, + ItemAccess: stubItemAccessChecker{}, + }) + handler.Transfers = registry + req := withEbookReaderRouteParams(newEbookReaderAuthRequest(http.MethodGet, "/ebooks/ebook-1/files/42/read"), "ebook-1", "42") + writer := &blockingEbookWriter{ + header: make(http.Header), + reached: make(chan struct{}), + release: make(chan struct{}), + } + done := make(chan struct{}) + go func() { + defer close(done) + handler.HandleReadFile(writer, req) + }() + + select { + case <-writer.reached: + case <-time.After(5 * time.Second): + t.Fatal("ebook pour did not reach the blocking writer") + } + active := registry.Snapshot() + if len(active) != 1 || active[0].Route != "ebook_read" || active[0].BytesServed == 0 { + t.Fatalf("active transfer = %+v, want metered ebook_read pour", active) + } + close(writer.release) + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("ebook pour did not finish") + } + if got := registry.Snapshot(); len(got) != 0 { + t.Fatalf("registry after completed pour = %+v, want empty", got) + } +} + +type revokingEbookResolver struct { + file *models.MediaFile + revoke func() +} + +func (r revokingEbookResolver) GetByID(context.Context, int) (*models.MediaFile, error) { + r.revoke() + return r.file, nil +} + +func TestEbookReaderRefusesRevocationLandingBeforePour(t *testing.T) { + filePath := writePlaybackTestMediaFile(t, "book.epub") + store := streamrevoke.New(streamrevoke.Options{}) + handler := NewEbookReaderHandler(&MediaFileAuthorizer{ + FileResolver: revokingEbookResolver{ + file: &models.MediaFile{ + ID: 42, ContentID: "ebook-1", FilePath: filePath, Container: "epub", BaseType: "ebook", + }, + revoke: func() { + if err := store.RevokeUser(context.Background(), 1, "test"); err != nil { + t.Fatal(err) + } + }, + }, + ItemAccess: stubItemAccessChecker{}, + }) + handler.Revocation = store + req := withEbookReaderRouteParams(newEbookReaderAuthRequest(http.MethodGet, "/ebooks/ebook-1/files/42/read"), "ebook-1", "42") + rr := httptest.NewRecorder() + + handler.HandleReadFile(rr, req) + + if rr.Code != http.StatusForbidden { + t.Fatalf("status = %d body = %q, want 403", rr.Code, rr.Body.String()) + } +} + +func TestEbookReaderRevocationCutsInFlightRealSocket(t *testing.T) { + filePath := writePlaybackTestMediaFile(t, "large-cut.epub") + if err := os.Truncate(filePath, 128<<20); err != nil { + t.Fatal(err) + } + store := streamrevoke.New(streamrevoke.Options{WatchInterval: 50 * time.Millisecond}) + handler := NewEbookReaderHandler(&MediaFileAuthorizer{ + FileResolver: stubMediaFileResolver{file: &models.MediaFile{ + ID: 42, ContentID: "ebook-1", FilePath: filePath, Container: "epub", BaseType: "ebook", + }}, + ItemAccess: stubItemAccessChecker{}, + }) + handler.Revocation = store + router := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := apimw.SetClaims(r.Context(), &auth.Claims{ + UserID: 1, Role: "user", TokenType: auth.TokenTypeAccess, + }) + ctx = apimw.SetProfileID(ctx, "profile-1") + routeCtx := chi.NewRouteContext() + routeCtx.URLParams.Add("content_id", "ebook-1") + routeCtx.URLParams.Add("file_id", "42") + ctx = context.WithValue(ctx, chi.RouteCtxKey, routeCtx) + handler.HandleReadFile(w, r.WithContext(ctx)) + }) + server := httptest.NewServer(router) + defer server.Close() + + response, err := http.Get(server.URL + "/ebooks/ebook-1/files/42/read") + if err != nil { + if response != nil { + _ = response.Body.Close() + } + t.Fatal(err) + } + defer func() { _ = response.Body.Close() }() + if err := store.RevokeUser(context.Background(), 1, "test"); err != nil { + t.Fatal(err) + } + readDone := make(chan error, 1) + go func() { + _, readErr := io.Copy(io.Discard, response.Body) + readDone <- readErr + }() + select { + case readErr := <-readDone: + if readErr == nil { + t.Fatal("client reached EOF; ebook revocation did not hang up the socket") + } + case <-time.After(8 * time.Second): + t.Fatal("revocation did not cut the in-flight ebook pour") + } +} + func TestEbookReaderServesContainerMimeForMismatchedExtension(t *testing.T) { filePath := writePlaybackTestMediaFile(t, "book.txt") if err := os.WriteFile(filePath, []byte("not really"), 0o644); err != nil { diff --git a/internal/api/handlers/playback.go b/internal/api/handlers/playback.go index 232a675a..f9f154c6 100644 --- a/internal/api/handlers/playback.go +++ b/internal/api/handlers/playback.go @@ -3889,9 +3889,26 @@ func (h *PlaybackHandler) proxyToTranscodeNode(w http.ResponseWriter, r *http.Re } // Proxied transcode output can stream past the server's absolute // WriteTimeout; roll the write deadline with progress instead. - sw := httpstream.NewRollingDeadlineWriter(w) + dst := w + var metered *playback.SessionMeteredWriter + if !strings.HasSuffix(path, ".m3u8") { + if err := h.sessionMgr.BeginTransport(sessionID); err == nil { + defer func() { _ = h.sessionMgr.EndTransport(sessionID) }() + } + recorder, _ := h.sessionMgr.(playback.ServedBytesRecorder) + metered = playback.NewSessionMeteredWriter(w, recorder, sessionID) + defer func() { _ = metered.Close() }() + dst = metered + } + sw := httpstream.NewRollingDeadlineWriterCtx(r.Context(), dst) sw.WriteHeader(resp.StatusCode) - io.Copy(sw, resp.Body) + if _, copyErr := io.Copy(sw, resp.Body); copyErr != nil { + slog.DebugContext(r.Context(), "copy transcode node response", + "component", "api", + "playback_session_id", sessionID, + "error", copyErr, + ) + } } // maybeStartThrottler reads throttle settings and starts the throttler if enabled. diff --git a/internal/api/handlers/playback_test.go b/internal/api/handlers/playback_test.go index d1100d3f..e32e2b09 100644 --- a/internal/api/handlers/playback_test.go +++ b/internal/api/handlers/playback_test.go @@ -1,17 +1,21 @@ package handlers import ( + "bytes" "context" "database/sql" "encoding/json" "errors" "fmt" + "io" + "log/slog" "math" "net/http" "net/http/httptest" "net/url" "os" "path/filepath" + "slices" "sort" "strings" "sync" @@ -35,6 +39,32 @@ import ( "github.com/Silo-Server/silo-server/internal/watchsync" ) +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +type recordingProxySessionManager struct { + *playback.SessionManager + events []string +} + +func (m *recordingProxySessionManager) BeginTransport(sessionID string) error { + m.events = append(m.events, "begin") + return m.SessionManager.BeginTransport(sessionID) +} + +func (m *recordingProxySessionManager) EndTransport(sessionID string) error { + m.events = append(m.events, "end") + return m.SessionManager.EndTransport(sessionID) +} + +func (m *recordingProxySessionManager) AddServedBytes(sessionID string, n int64) error { + m.events = append(m.events, fmt.Sprintf("bytes:%d", n)) + return m.SessionManager.AddServedBytes(sessionID, n) +} + type testUserStoreProvider struct { store userstore.UserStore } @@ -1981,6 +2011,96 @@ func TestHandleChangeAudioTrack_RemoteTranscodeNodeUnreachableSurfacesError(t *t } } +func TestProxyToTranscodeNodeMetersOnlyMediaBodiesAndFlushesBeforeTransportEnd(t *testing.T) { + const body = "segment-bytes" + previousTransport := http.DefaultClient.Transport + http.DefaultClient.Transport = roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewBufferString(body)), + }, nil + }) + t.Cleanup(func() { http.DefaultClient.Transport = previousTransport }) + + for _, tc := range []struct { + name string + path string + wantBytes int64 + wantEvents []string + }{ + { + name: "segment", + path: "/transcode/s/segment/part.m4s", + wantBytes: int64(len(body)), + wantEvents: []string{"begin", fmt.Sprintf("bytes:%d", len(body)), "end"}, + }, + {name: "playlist body", path: "/transcode/s/master.m3u8"}, + } { + t.Run(tc.name, func(t *testing.T) { + manager := &recordingProxySessionManager{SessionManager: playback.NewSessionManager(0, 0)} + session, err := manager.StartSession(1, "profile-1", 42, playback.PlayTranscode, false) + if err != nil { + t.Fatal(err) + } + handler := NewPlaybackHandler(manager) + handler.JWTSecret = "test-secret" + req := httptest.NewRequest(http.MethodGet, "/stream/"+session.ID, nil) + routeCtx := chi.NewRouteContext() + routeCtx.URLParams.Add("session_id", session.ID) + req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, routeCtx)) + rec := httptest.NewRecorder() + + handler.proxyToTranscodeNode(rec, req, "http://transcode.test", tc.path) + + got, err := manager.GetSession(session.ID) + if err != nil { + t.Fatal(err) + } + if got.BytesServed != tc.wantBytes { + t.Fatalf("bytes served = %d, want %d", got.BytesServed, tc.wantBytes) + } + if !slices.Equal(manager.events, tc.wantEvents) { + t.Fatalf("events = %v, want %v", manager.events, tc.wantEvents) + } + }) + } +} + +type failingCopyBody struct{} + +func (failingCopyBody) Read([]byte) (int, error) { return 0, errors.New("copy failed") } +func (failingCopyBody) Close() error { return nil } + +func TestProxyToTranscodeNodeLogsCopyFailure(t *testing.T) { + previousTransport := http.DefaultClient.Transport + http.DefaultClient.Transport = roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: failingCopyBody{}, + }, nil + }) + t.Cleanup(func() { http.DefaultClient.Transport = previousTransport }) + + var logs bytes.Buffer + previousLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(previousLogger) }) + + handler := NewPlaybackHandler(failingSessionManager{}) + req := httptest.NewRequest(http.MethodGet, "/stream/session-1", nil) + routeCtx := chi.NewRouteContext() + routeCtx.URLParams.Add("session_id", "session-1") + req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, routeCtx)) + + handler.proxyToTranscodeNode(httptest.NewRecorder(), req, "http://transcode.test", "/transcode/session-1/segment/part.m4s") + + if !strings.Contains(logs.String(), "copy transcode node response") || !strings.Contains(logs.String(), "copy failed") { + t.Fatalf("copy failure log missing: %s", logs.String()) + } +} + func TestHandleStartTranscode_LocalPathPropagatesSelectedAudioTrack(t *testing.T) { sessionMgr := playback.NewSessionManager(0, 0) filePath := writePlaybackTestMediaFile(t, "movie.mkv") diff --git a/internal/api/handlers/stream.go b/internal/api/handlers/stream.go index d6541060..dbcce7e7 100644 --- a/internal/api/handlers/stream.go +++ b/internal/api/handlers/stream.go @@ -224,6 +224,20 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusNotFound, "not_found", "Media file not found") return } + beginTransport := func() (http.ResponseWriter, func()) { + transportStarted := false + if err := h.sessionMgr.BeginTransport(sessionID); err == nil { + transportStarted = true + } + recorder, _ := h.sessionMgr.(playback.ServedBytesRecorder) + metered := playback.NewSessionMeteredWriter(w, recorder, sessionID) + return metered, func() { + _ = metered.Close() + if transportStarted { + _ = h.sessionMgr.EndTransport(sessionID) + } + } + } externalCount := len(file.ExternalSubtitles) if trackIndex < externalCount { @@ -237,7 +251,9 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { "Failed to load external subtitle") return } - playback.ServeSubtitle(w, data, "ass") + streamWriter, finish := beginTransport() + defer finish() + playback.ServeSubtitle(streamWriter, data, "ass") return } @@ -247,7 +263,9 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { "Failed to load external subtitle") return } - playback.ServeSubtitle(w, vttData, "vtt") + streamWriter, finish := beginTransport() + defer finish() + playback.ServeSubtitle(streamWriter, vttData, "vtt") return } @@ -270,7 +288,9 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { // demuxed, so the first byte lands within ~1s even on network // storage. Works identically for direct-play, remux, and // transcode because it doesn't depend on any other ffmpeg. - h.streamEmbeddedSubtitle(w, r, file, embeddedIndex, session, requestedFormat) + streamWriter, finish := beginTransport() + defer finish() + h.streamEmbeddedSubtitle(streamWriter, r, file, embeddedIndex, session, requestedFormat) return } @@ -302,13 +322,17 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { // Serve ASS/SSA downloaded subtitles as raw data. if playback.IsASS(string(dl.Format)) && requestedFormat != "vtt" { - playback.ServeSubtitle(w, data, "ass") + streamWriter, finish := beginTransport() + defer finish() + playback.ServeSubtitle(streamWriter, data, "ass") return } // If the subtitle is already VTT, serve directly. if dl.Format == subtitles.FormatVTT { - playback.ServeSubtitle(w, data, "vtt") + streamWriter, finish := beginTransport() + defer finish() + playback.ServeSubtitle(streamWriter, data, "vtt") return } @@ -318,7 +342,9 @@ func (h *StreamHandler) HandleSubtitle(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusInternalServerError, "convert_error", "Failed to convert subtitle") return } - playback.ServeSubtitle(w, vttData, "vtt") + streamWriter, finish := beginTransport() + defer finish() + playback.ServeSubtitle(streamWriter, vttData, "vtt") return } } @@ -426,6 +452,14 @@ func (h *StreamHandler) HandleSubtitleFonts(w http.ResponseWriter, r *http.Reque return } + if err := h.sessionMgr.BeginTransport(sessionID); err == nil { + defer func() { _ = h.sessionMgr.EndTransport(sessionID) }() + } + recorder, _ := h.sessionMgr.(playback.ServedBytesRecorder) + metered := playback.NewSessionMeteredWriter(w, recorder, sessionID) + defer func() { _ = metered.Close() }() + w = metered + w.Header().Set("Content-Type", "application/json") w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("Cache-Control", "no-store") diff --git a/internal/api/router.go b/internal/api/router.go index db8cd1d5..167af795 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -36,6 +36,7 @@ import ( "github.com/Silo-Server/silo-server/internal/downloads" evt "github.com/Silo-Server/silo-server/internal/events" "github.com/Silo-Server/silo-server/internal/historyimport" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/intromarkers" "github.com/Silo-Server/silo-server/internal/invitations" "github.com/Silo-Server/silo-server/internal/libraryingest" @@ -687,6 +688,8 @@ func NewRouter(deps Dependencies) chi.Router { if conv := buildEbookConversion(deps, settingsRepo); conv != nil { ebookReaderHandler.Conversion = conv } + ebookReaderHandler.Transfers = deps.TransferRegistry + ebookReaderHandler.Revocation = deps.RevocationStore } catalogResourceHandler = handlers.NewCatalogResourceHandler(itemsHandler) catalogHandler = handlers.NewCatalogHandler( @@ -2540,8 +2543,8 @@ func NewRouter(deps Dependencies) chi.Router { if streamHandler != nil { r.Get("/stream/{session_id}", guardRevocationCut(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleStream)) r.Head("/stream/{session_id}", guardRevocationCut(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleStream)) - r.Get("/stream/{session_id}/subtitles/{track}", guardRevocation(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleSubtitle)) - r.Get("/stream/{session_id}/subtitles/{track}/fonts", guardRevocation(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleSubtitleFonts)) + r.Get("/stream/{session_id}/subtitles/{track}", guardRevocationCut(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleSubtitle)) + r.Get("/stream/{session_id}/subtitles/{track}/fonts", guardRevocationCut(deps.RevocationStore, configJWTSecret(deps), streamHandler.HandleSubtitleFonts)) } // Download routes. @@ -3566,10 +3569,12 @@ func guardRevocationCut(store *streamrevoke.Store, secret string, next http.Hand return func(w http.ResponseWriter, r *http.Request) { sessionID := chi.URLParam(r, "session_id") userID, startedAt := streamRequestIdentity(r, secret) + latch := &httpstream.CutLatch{} + r = r.WithContext(httpstream.WithCutLatch(r.Context(), latch)) if store.Refuse(w, sessionID, userID, startedAt) { return } - stop := store.WatchAndCut(w, sessionID, userID, startedAt) + stop := store.WatchAndCutContext(r.Context(), w, sessionID, userID, startedAt) defer stop() next(w, r) } diff --git a/internal/audiobooks/abs/access_log.go b/internal/audiobooks/abs/access_log.go index 8416732f..d3111c4c 100644 --- a/internal/audiobooks/abs/access_log.go +++ b/internal/audiobooks/abs/access_log.go @@ -87,6 +87,14 @@ type statusRecorder struct { bytes int } +// Unwrap exposes the wrapped ResponseWriter to http.NewResponseController so +// write-deadline control — the revocation in-flight cut — reaches the real +// socket. Without this, SetWriteDeadline returns ErrNotSupported and a +// RevokeUser cannot hang up an in-flight multi-GB audiobook pour. +func (s *statusRecorder) Unwrap() http.ResponseWriter { + return s.ResponseWriter +} + func (s *statusRecorder) WriteHeader(code int) { s.status = code s.ResponseWriter.WriteHeader(code) diff --git a/internal/audiobooks/abs/file_handler.go b/internal/audiobooks/abs/file_handler.go index b6f4f193..35c3ec16 100644 --- a/internal/audiobooks/abs/file_handler.go +++ b/internal/audiobooks/abs/file_handler.go @@ -13,6 +13,7 @@ import ( "github.com/go-chi/chi/v5" "github.com/google/uuid" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/playback" "github.com/Silo-Server/silo-server/internal/transfers" ) @@ -56,11 +57,10 @@ func trackInoFor(contentID string, fileIdx int) string { // - Set Content-Disposition: attachment on /download paths to encourage // browser save-to-disk / mobile offline-save behaviour. // -// This is a download-class route: like native and Jellyfin-compatible -// downloads it is exempt from the live-stream cap and never enters -// streammonitor. Its active pour is exposed through the separate process-local -// transfer registry. It is not yet covered by a download quota and remains -// revocable before and during the file pour. +// Both forms serve real playback bytes, but per A4 they are deliberately +// exempt from the live-stream cap and absent from streammonitor. Each active +// pour is still metered, registered in the process-local transfer registry, +// and revocable before and during transfer. func (h *Handler) handleFileStream(w http.ResponseWriter, r *http.Request) { a, ok := absAuthFrom(r) if !ok || a.UserID == "" { @@ -72,6 +72,7 @@ func (h *Handler) handleFileStream(w http.ResponseWriter, r *http.Request) { http.Error(w, "unauthorized", http.StatusUnauthorized) return } + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) if h.deps.Revocation != nil && h.deps.Revocation.Refuse(w, "", userID, a.IssuedAt) { return } @@ -156,7 +157,7 @@ func (h *Handler) handleFileStream(w http.ResponseWriter, r *http.Request) { defer func() { _ = metered.Close() }() if h.deps.Revocation != nil { - stop := h.deps.Revocation.WatchAndCut(metered, "", userID, a.IssuedAt) + stop := h.deps.Revocation.WatchAndCutContext(r.Context(), metered, "", userID, a.IssuedAt) defer stop() } if err := playback.ServeDirectPlay(metered, r, mediaFile.FilePath); err != nil { @@ -221,6 +222,7 @@ func (h *Handler) handlePublicTrack(w http.ResponseWriter, r *http.Request) { http.Error(w, "session owner invalid", http.StatusForbidden) return } + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) if h.deps.Revocation != nil && h.deps.Revocation.Refuse(w, sid, userID, sess.StartedAt) { return } @@ -251,7 +253,7 @@ func (h *Handler) handlePublicTrack(w http.ResponseWriter, r *http.Request) { recorder, _ := h.deps.NativeSessions.(playback.ServedBytesRecorder) metered := playback.NewSessionMeteredWriter(w, recorder, sid) defer func() { _ = metered.Close() }() - stop := h.deps.Revocation.WatchAndCut(metered, sid, userID, sess.StartedAt) + stop := h.deps.Revocation.WatchAndCutContext(r.Context(), metered, sid, userID, sess.StartedAt) defer stop() _ = playback.ServeDirectPlay(metered, r, mediaFile.FilePath) } diff --git a/internal/audiobooks/abs/revocation_test.go b/internal/audiobooks/abs/revocation_test.go index c2d937af..99b8c885 100644 --- a/internal/audiobooks/abs/revocation_test.go +++ b/internal/audiobooks/abs/revocation_test.go @@ -2,17 +2,91 @@ package abs import ( "context" + "io" "net/http" "net/http/httptest" + "sync" "testing" "time" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/models" "github.com/Silo-Server/silo-server/internal/streamrevoke" "github.com/go-chi/chi/v5" "github.com/golang-jwt/jwt/v5" ) +var _ interface{ Unwrap() http.ResponseWriter } = (*statusRecorder)(nil) + +func TestStatusRecorderLetsResponseControllerReachSocketDeadline(t *testing.T) { + base := &deadlineResponseRecorder{ResponseRecorder: httptest.NewRecorder()} + wrapped := &statusRecorder{ResponseWriter: base} + if err := http.NewResponseController(wrapped).SetWriteDeadline(time.Now()); err != nil { + t.Fatalf("SetWriteDeadline through statusRecorder: %v", err) + } + if base.deadlines != 1 { + t.Fatalf("underlying deadline calls = %d, want 1", base.deadlines) + } +} + +func TestMountedAccessLogMiddlewareLetsRevocationCutRealSocket(t *testing.T) { + store := streamrevoke.New(streamrevoke.Options{WatchInterval: 50 * time.Millisecond}) + startedAt := time.Now() + started := make(chan struct{}) + var startedOnce sync.Once + h := &Handler{} + router := chi.NewRouter() + router.Use(h.accessLog) + router.Get("/pour", func(w http.ResponseWriter, r *http.Request) { + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) + stop := store.WatchAndCutContext(r.Context(), w, "abs-mounted", 1, startedAt) + defer stop() + sw := httpstream.NewRollingDeadlineWriterCtx(r.Context(), w) + sw.WriteHeader(http.StatusOK) + chunk := make([]byte, 32<<10) + for { + if _, err := sw.Write(chunk); err != nil { + return + } + sw.Flush() + startedOnce.Do(func() { close(started) }) + time.Sleep(10 * time.Millisecond) + } + }) + server := httptest.NewServer(router) + defer server.Close() + + response, err := http.Get(server.URL + "/pour") + if err != nil { + if response != nil { + _ = response.Body.Close() + } + t.Fatal(err) + } + defer func() { _ = response.Body.Close() }() + select { + case <-started: + case <-time.After(3 * time.Second): + t.Fatal("mounted ABS pour did not start") + } + if err := store.RevokeSession(context.Background(), "abs-mounted", "test"); err != nil { + t.Fatal(err) + } + readDone := make(chan error, 1) + go func() { + _, readErr := io.Copy(io.Discard, response.Body) + readDone <- readErr + }() + select { + case readErr := <-readDone: + if readErr == nil { + t.Fatal("client reached EOF; revocation did not hang up the socket") + } + case <-time.After(8 * time.Second): + t.Fatal("revocation did not cut the mounted ABS pour") + } +} + func TestBearerAuthCarriesJWTTimestampForUserCutoff(t *testing.T) { const accessTokenType = "access" secret := []byte("test-secret-32-bytes-aaaaaaaaaaaaa") diff --git a/internal/audiobooks/abs/rss_feeds_handler.go b/internal/audiobooks/abs/rss_feeds_handler.go index 54f02b3b..5d98ec46 100644 --- a/internal/audiobooks/abs/rss_feeds_handler.go +++ b/internal/audiobooks/abs/rss_feeds_handler.go @@ -16,6 +16,7 @@ import ( "github.com/google/uuid" "github.com/oklog/ulid/v2" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/models" "github.com/Silo-Server/silo-server/internal/playback" "github.com/Silo-Server/silo-server/internal/transfers" @@ -232,6 +233,7 @@ func (h *Handler) handlePublicFeedFile(w http.ResponseWriter, r *http.Request) { http.Error(w, "feed not found", http.StatusNotFound) return } + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) if h.deps.Revocation != nil && h.deps.Revocation.Refuse(w, "", userID, f.CreatedAt) { return } @@ -266,7 +268,7 @@ func (h *Handler) handlePublicFeedFile(w http.ResponseWriter, r *http.Request) { defer func() { _ = metered.Close() }() if h.deps.Revocation != nil { - stop := h.deps.Revocation.WatchAndCut(metered, "", userID, f.CreatedAt) + stop := h.deps.Revocation.WatchAndCutContext(r.Context(), metered, "", userID, f.CreatedAt) defer stop() } http.ServeFile(metered, r, mf.FilePath) diff --git a/internal/jellycompat/streams.go b/internal/jellycompat/streams.go index 0a690599..dc041e77 100644 --- a/internal/jellycompat/streams.go +++ b/internal/jellycompat/streams.go @@ -22,6 +22,7 @@ import ( "github.com/go-chi/chi/v5" "github.com/Silo-Server/silo-server/internal/clientip" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/models" "github.com/Silo-Server/silo-server/internal/nodepool" "github.com/Silo-Server/silo-server/internal/playback" @@ -56,6 +57,7 @@ func (h *PlaybackHandler) HandleVideoStream(w http.ResponseWriter, r *http.Reque writeError(w, http.StatusUnauthorized, "Unauthorized", "Missing authentication token") return } + startedAt := time.Now() routeID := chiURLParam(r, "id") mediaSourceID := firstNonEmpty(r.URL.Query().Get("mediaSourceId"), r.URL.Query().Get("MediaSourceId")) @@ -97,7 +99,7 @@ func (h *PlaybackHandler) HandleVideoStream(w http.ResponseWriter, r *http.Reque // stream URL. (User-kind kills don't need this: they match by user id either // side of the ensure.) if h.Revocation != nil && playSession.UpstreamSessionID != "" && - h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, time.Now()) { + h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, startedAt) { return } @@ -114,7 +116,7 @@ func (h *PlaybackHandler) HandleVideoStream(w http.ResponseWriter, r *http.Reque // request only gets here on a live compat login, and user revocation deletes // all compat logins, so reaching this line post-revocation means fresh auth. if h.Revocation != nil && playSession.UpstreamSessionID != "" && - h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, time.Now()) { + h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, startedAt) { return } @@ -155,8 +157,9 @@ func (h *PlaybackHandler) HandleVideoStream(w http.ResponseWriter, r *http.Reque // In-flight cut: hang up a long local direct-play/remux pour the moment it is // revoked. The Refuse check above only covers new/reconnect requests; a single // long GET needs the connection cut to stop mid-stream. + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) if h.Revocation != nil && playSession.UpstreamSessionID != "" { - stop := h.Revocation.WatchAndCut(w, playSession.UpstreamSessionID, session.StreamAppUserID, time.Now()) + stop := h.Revocation.WatchAndCutContext(r.Context(), w, playSession.UpstreamSessionID, session.StreamAppUserID, startedAt) defer stop() } recorder, _ := h.sessionMgr.(playback.ServedBytesRecorder) @@ -258,7 +261,8 @@ func (h *PlaybackHandler) HandleDownload(w http.ResponseWriter, r *http.Request) // the entry time predates any future revocation (the user-kill cutoff), so a // revocation issued mid-transfer hangs this connection up. if h.Revocation != nil { - stop := h.Revocation.WatchAndCut(metered, "", session.StreamAppUserID, transfer.StartedAt) + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) + stop := h.Revocation.WatchAndCutContext(r.Context(), metered, "", session.StreamAppUserID, transfer.StartedAt) defer stop() } @@ -680,13 +684,19 @@ func (h *PlaybackHandler) HandleSubtitleStream(w http.ResponseWriter, r *http.Re return } + startedAt := time.Now() // Kill switch: a revoked session must not keep pulling subtitle bytes (or - // triggering server-side ffmpeg subtitle extraction below). Mirrors the - // guarded native subtitle routes. + // triggering server-side ffmpeg subtitle extraction below). Compat + // extraction is buffered, so a socket cut stops response delivery but does + // not interrupt extraction already in progress. if h.Revocation != nil && playSession != nil && playSession.UpstreamSessionID != "" && - h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, time.Now()) { + h.Revocation.Refuse(w, playSession.UpstreamSessionID, session.StreamAppUserID, startedAt) { return } + if h.Revocation != nil && playSession != nil && playSession.UpstreamSessionID != "" { + stop := h.Revocation.WatchAndCutContext(r.Context(), w, playSession.UpstreamSessionID, session.StreamAppUserID, startedAt) + defer stop() + } if h.fileResolver == nil { writeError(w, http.StatusInternalServerError, "ServerError", "File resolver not available") diff --git a/internal/proxy/remove_tracked_test.go b/internal/proxy/remove_tracked_test.go new file mode 100644 index 00000000..ff03366d --- /dev/null +++ b/internal/proxy/remove_tracked_test.go @@ -0,0 +1,63 @@ +package proxy + +import ( + "context" + "net" + "testing" + + "github.com/Silo-Server/silo-server/internal/nodesessions" + "github.com/redis/go-redis/v9" +) + +type deleteCaptureHook struct { + called bool + contextErr error +} + +func (h *deleteCaptureHook) DialHook(next redis.DialHook) redis.DialHook { + return next +} + +func (h *deleteCaptureHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if cmd.Name() == "del" { + h.called = true + h.contextErr = ctx.Err() + return nil + } + return next(ctx, cmd) + } +} + +func (h *deleteCaptureHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return next +} + +func TestRemoveTrackedUsesLiveContextAfterRequestCancellation(t *testing.T) { + rdb := redis.NewClient(&redis.Options{ + Dialer: func(context.Context, string, string) (net.Conn, error) { + t.Fatal("DEL should be intercepted before dialing") + return nil, nil + }, + }) + t.Cleanup(func() { _ = rdb.Close() }) + hook := &deleteCaptureHook{} + rdb.AddHook(hook) + server := &Server{tracker: nodesessions.NewTracker(rdb, "http://node", "node", "proxy")} + + requestCtx, cancel := context.WithCancel(context.Background()) + cancel() + if requestCtx.Err() == nil { + t.Fatal("request context was not canceled") + } + func() { + defer server.removeTracked("session-1") + }() + + if !hook.called { + t.Fatal("Redis DEL was not attempted") + } + if hook.contextErr != nil { + t.Fatalf("DEL context error = %v, want live bounded cleanup context", hook.contextErr) + } +} diff --git a/internal/proxy/revocation_test.go b/internal/proxy/revocation_test.go new file mode 100644 index 00000000..cb9bf1ee --- /dev/null +++ b/internal/proxy/revocation_test.go @@ -0,0 +1,60 @@ +package proxy + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Silo-Server/silo-server/internal/config" + "github.com/Silo-Server/silo-server/internal/httpstream" + "github.com/Silo-Server/silo-server/internal/nodeconfig" + "github.com/Silo-Server/silo-server/internal/nodesessions" + "github.com/Silo-Server/silo-server/internal/streamtoken" +) + +type latchCapturingRevocationStore struct { + latch *httpstream.CutLatch +} + +func (s *latchCapturingRevocationStore) IsRevoked(string, int, time.Time) bool { + return false +} + +func (s *latchCapturingRevocationStore) Refuse(http.ResponseWriter, string, int, time.Time) bool { + return false +} + +func (s *latchCapturingRevocationStore) WatchAndCutContext(ctx context.Context, _ http.ResponseWriter, _ string, _ int, _ time.Time) func() { + s.latch = httpstream.CutLatchFrom(ctx) + return func() {} +} + +func TestHandleRemuxCarriesCutLatchToRevocationWatcher(t *testing.T) { + const secret = "proxy-revocation-test-secret" + watcher := nodeconfig.NewWatcher(nil, nil, nil, nodeconfig.BootstrapOverrides{}) + cfg := &config.Config{} + cfg.Auth.JWTSecret = secret + watcher.SetConfigForTest(cfg) + + server := NewServer(watcher, nodesessions.NewTracker(nil, "http://proxy", "proxy", "proxy")) + revocation := &latchCapturingRevocationStore{} + server.SetRevocationStore(revocation) + + token, err := streamtoken.Sign(streamtoken.Claims{ + SessionID: "remux-latch", + MediaPath: t.TempDir() + "/missing.mkv", + }, secret, time.Minute) + if err != nil { + t.Fatal(err) + } + req := httptest.NewRequest(http.MethodGet, "/stream/remux/"+token, nil) + rr := httptest.NewRecorder() + + server.Handler().ServeHTTP(rr, req) + + if revocation.latch == nil { + t.Fatal("remux revocation watcher request context has no cut latch") + } +} diff --git a/internal/proxy/server.go b/internal/proxy/server.go index 09a07b56..709e53d6 100644 --- a/internal/proxy/server.go +++ b/internal/proxy/server.go @@ -1,6 +1,7 @@ package proxy import ( + "context" "encoding/json" "fmt" "io" @@ -14,6 +15,7 @@ import ( "github.com/go-chi/chi/v5" "github.com/go-chi/cors" + "github.com/Silo-Server/silo-server/internal/httpstream" "github.com/Silo-Server/silo-server/internal/nodeconfig" "github.com/Silo-Server/silo-server/internal/nodesessions" "github.com/Silo-Server/silo-server/internal/playback" @@ -26,7 +28,7 @@ import ( type revocationStore interface { IsRevoked(sessionID string, userID int, startedAt time.Time) bool Refuse(w http.ResponseWriter, sessionID string, userID int, startedAt time.Time) bool - WatchAndCut(w http.ResponseWriter, sessionID string, userID int, startedAt time.Time) func() + WatchAndCutContext(ctx context.Context, w http.ResponseWriter, sessionID string, userID int, startedAt time.Time) func() } // Server is the HTTP handler for proxy mode. @@ -172,14 +174,25 @@ func (s *Server) verifyToken(w http.ResponseWriter, r *http.Request) *streamtoke // cutOnRevocation watches a long-lived pour (direct play / remux) and hangs up // the socket the moment the session/user is revoked, via the shared store helper. // Returns a stop func to cancel the watcher when the request finishes normally. -func (s *Server) cutOnRevocation(w http.ResponseWriter, claims *streamtoken.Claims) func() { +func (s *Server) cutOnRevocation(ctx context.Context, w http.ResponseWriter, claims *streamtoken.Claims) func() { if s.revocation == nil { return func() {} } - return s.revocation.WatchAndCut(w, claims.SessionID, claims.UserID, claims.IssuedTime()) + return s.revocation.WatchAndCutContext(ctx, w, claims.SessionID, claims.UserID, claims.IssuedTime()) +} + +// removeTracked deletes the edge record with a bounded background context. The +// request context is already canceled by the time this defer runs on a client +// disconnect, which would skip the Redis DEL and leave a phantom session until +// TTL — a false over-cap window. +func (s *Server) removeTracked(sessionID string) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + s.tracker.Remove(ctx, sessionID) } func (s *Server) handleDirectPlay(w http.ResponseWriter, r *http.Request) { + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) claims := s.verifyToken(w, r) if claims == nil { return @@ -187,12 +200,12 @@ func (s *Server) handleDirectPlay(w http.ResponseWriter, r *http.Request) { info := sessionInfo(s.tracker, claims, "direct_play", edgeClientIP(r)) s.tracker.Track(r.Context(), info) - defer s.tracker.Remove(r.Context(), claims.SessionID) + defer s.removeTracked(claims.SessionID) sw := &sessionByteWriter{ResponseWriter: w, tracker: s.tracker, sessionID: claims.SessionID} defer sw.flush() - stop := s.cutOnRevocation(sw, claims) + stop := s.cutOnRevocation(r.Context(), sw, claims) defer stop() http.ServeFile(sw, r, claims.MediaPath) @@ -272,6 +285,7 @@ func (w *sessionByteWriter) Unwrap() http.ResponseWriter { } func (s *Server) handleRemux(w http.ResponseWriter, r *http.Request) { + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) claims := s.verifyToken(w, r) if claims == nil { return @@ -279,12 +293,12 @@ func (s *Server) handleRemux(w http.ResponseWriter, r *http.Request) { info := sessionInfo(s.tracker, claims, "remux", edgeClientIP(r)) s.tracker.Track(r.Context(), info) - defer s.tracker.Remove(r.Context(), claims.SessionID) + defer s.removeTracked(claims.SessionID) sw := &sessionByteWriter{ResponseWriter: w, tracker: s.tracker, sessionID: claims.SessionID} defer sw.flush() - stop := s.cutOnRevocation(sw, claims) + stop := s.cutOnRevocation(r.Context(), sw, claims) defer stop() seekSeconds := 0.0 @@ -369,10 +383,21 @@ func edgeClientIP(r *http.Request) string { } func (s *Server) handleSubtitle(w http.ResponseWriter, r *http.Request) { + r = r.WithContext(httpstream.WithCutLatch(r.Context(), &httpstream.CutLatch{})) claims := s.verifyToken(w, r) if claims == nil { return } + // Attribute subtitle bytes only when the session already has a tracker + // record. Do not take lifecycle ownership here: Track/Remove per subtitle + // request can delete a concurrently active media request until Batch 2 adds + // reference-counted edge tracking. + metered := &sessionByteWriter{ResponseWriter: w, tracker: s.tracker, sessionID: claims.SessionID} + defer metered.flush() + stop := s.cutOnRevocation(r.Context(), metered, claims) + defer stop() + w = metered + cfg := s.watcher.Config() trackParam := chi.URLParam(r, "track") trackIndex, requestedFormat, err := playback.ParseSubtitleTrackParam(trackParam) diff --git a/internal/streammonitor/monitor.go b/internal/streammonitor/monitor.go index 2686977d..64e14640 100644 --- a/internal/streammonitor/monitor.go +++ b/internal/streammonitor/monitor.go @@ -227,6 +227,8 @@ func mergeStreams(streams []LiveStream) []LiveStream { if winner.Position == 0 { winner.Position = other.Position } + // These are observers of one pour, not independent byte sources. + winner.BytesServed = max(winner.BytesServed, other.BytesServed) bySession[st.SessionID] = winner } out := make([]LiveStream, 0, len(bySession)) @@ -278,6 +280,7 @@ func DedupeSessionInfos(infos []nodesessions.SessionInfo) []nodesessions.Session if winner.Position == 0 { winner.Position = other.Position } + winner.BytesServed = max(winner.BytesServed, other.BytesServed) bySession[in.SessionID] = winner } out := make([]nodesessions.SessionInfo, 0, len(bySession)) diff --git a/internal/streammonitor/monitor_test.go b/internal/streammonitor/monitor_test.go index a8c2748a..8ae09a29 100644 --- a/internal/streammonitor/monitor_test.go +++ b/internal/streammonitor/monitor_test.go @@ -213,6 +213,27 @@ func TestMergeStreamsBackfillsAttribution(t *testing.T) { } } +func TestMergeStreamsKeepsLargestObservedByteTotal(t *testing.T) { + for _, tc := range []struct { + name string + older, newer int64 + want int64 + }{ + {name: "freshest has zero", older: 8192, newer: 0, want: 8192}, + {name: "stale has zero", older: 0, newer: 8192, want: 8192}, + } { + t.Run(tc.name, func(t *testing.T) { + out := mergeStreams([]LiveStream{ + {SessionID: "s", BytesServed: tc.older, LastServedAt: time.Unix(100, 0)}, + {SessionID: "s", BytesServed: tc.newer, LastServedAt: time.Unix(200, 0)}, + }) + if len(out) != 1 || out[0].BytesServed != tc.want { + t.Fatalf("merge = %+v, want bytes %d", out, tc.want) + } + }) + } +} + func TestFuncSourceNilFn(t *testing.T) { snap, err := NewFuncSource(nil).Snapshot(context.Background()) if err != nil { @@ -274,3 +295,26 @@ func TestDedupeSessionInfos(t *testing.T) { t.Errorf("s2 mangled by dedupe: %+v", s2) } } + +func TestDedupeSessionInfosKeepsLargestObservedByteTotal(t *testing.T) { + older := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339Nano) + newer := time.Now().UTC().Format(time.RFC3339Nano) + for _, tc := range []struct { + name string + older, newer int64 + want int64 + }{ + {name: "freshest has zero", older: 4096, newer: 0, want: 4096}, + {name: "stale has zero", older: 0, newer: 4096, want: 4096}, + } { + t.Run(tc.name, func(t *testing.T) { + out := DedupeSessionInfos([]nodesessions.SessionInfo{ + {SessionID: "s", BytesServed: tc.older, LastServedAt: older}, + {SessionID: "s", BytesServed: tc.newer, LastServedAt: newer}, + }) + if len(out) != 1 || out[0].BytesServed != tc.want { + t.Fatalf("dedupe = %+v, want bytes %d", out, tc.want) + } + }) + } +}