From b6d18a267b67918e438fc3920b45f5fc580d6cc1 Mon Sep 17 00:00:00 2001 From: sun Date: Tue, 15 Sep 2026 11:18:40 +0800 Subject: [PATCH] fix(hls): rebuild request extras on restore and adopt merge semantics in Patch - Parse the request extra in Setup so restored fetchers resend Referer, Cookie and other configured headers on resumed segment/key requests; the store/restore regression test round-trips the state and the task meta through JSON exactly like the engine storage does - Rebuild the resolved plan on Start when the resource survived but the plan was lost to a crash between the task save and the first checkpoint, keeping the existing output name - Patch now merges partial updates instead of replacing the request and options wholesale: URL-only patches keep the output path and renamed file, header-only patches keep the resolved plan and take effect on the next run, Labels merge per key, Proxy swaps rebuild the client, and nil arguments are no longer a panic. Editing a running task is rejected, consistent with the ed2k protocol - Cover retry-after-failure counter resets, running-task patch rejection and the missing-state self-heal with regression tests --- internal/protocol/hls/fetcher.go | 165 ++++++++++-- internal/protocol/hls/fetcher_test.go | 375 ++++++++++++++++++++++++++ 2 files changed, 523 insertions(+), 17 deletions(-) diff --git a/internal/protocol/hls/fetcher.go b/internal/protocol/hls/fetcher.go index 89eed251..9fc041b9 100644 --- a/internal/protocol/hls/fetcher.go +++ b/internal/protocol/hls/fetcher.go @@ -139,6 +139,25 @@ func (f *Fetcher) Setup(ctl *controller.Controller) { f.config = &config{} f.Ctl.GetConfig(f.config) f.config.normalize() + // Restored fetchers receive their request through Setup (the engine calls + // it after the Restore builder): parse the extra headers so resumed + // segment/key requests keep Referer/Cookie and other anti-hotlink headers. + f.mu.Lock() + f.reloadExtraLocked() + f.mu.Unlock() +} + +// reloadExtraLocked refreshes f.extra from the typed extra of meta.Req. +func (f *Fetcher) reloadExtraLocked() { + if f.meta == nil || f.meta.Req == nil { + return + } + if err := base.ParseReqExtra[ReqExtra](f.meta.Req); err != nil { + return + } + if e, ok := f.meta.Req.Extra.(*ReqExtra); ok { + f.extra = e + } } func (f *Fetcher) Meta() *fetcher.FetcherMeta { @@ -236,29 +255,119 @@ func (f *Fetcher) resolvePlanLocked(ctx context.Context, keepName string) (*fetc return state, outputName, nil } -func (f *Fetcher) Patch(req *base.Request, opts *base.Options) (err error) { +func (f *Fetcher) Patch(req *base.Request, opts *base.Options) error { f.mu.Lock() defer f.mu.Unlock() if f.activeRun != nil { return errors.New("task is running; pause it before changing its settings") } - urlChanged := f.meta.Req == nil || f.meta.Req.URL != req.URL - f.meta.Req = req - f.meta.Opts = opts - f.extra = &ReqExtra{} - if err = base.ParseReqExtra[ReqExtra](req); err != nil { - return + if req != nil { + cur := f.meta.Req + if cur == nil { + f.meta.Req = req + cur = req + } + if req.URL != "" && req.URL != cur.URL { + cur.URL = req.URL + // The playlist URL changed (e.g. refreshed signed link): drop the + // plan so the engine re-resolves on next start (it does when Res + // is nil). Only explicitly supplied fields are applied, so a + // header-only or rename-only Patch keeps the resolved plan and + // the output path. + f.dropPlanLocked() + } + if err := mergeReqExtra(cur, req); err != nil { + return err + } + if req.Proxy != nil && cur.Proxy != req.Proxy { + cur.Proxy = req.Proxy + // Transport settings changed: rebuild the client before the + // next run. + f.client = nil + } + if len(req.Labels) > 0 { + if cur.Labels == nil { + cur.Labels = make(map[string]string, len(req.Labels)) + } + for key, value := range req.Labels { + cur.Labels[key] = value + } + } + // SkipVerifyCert is a plain bool, so a partial update cannot tell an + // absent field from an explicit false: it is intentionally not merged + // (the HTTP protocol Patch ignores it as well). } - if e, ok := req.Extra.(*ReqExtra); ok { - f.extra = e + if opts != nil { + if f.meta.Opts == nil { + f.meta.Opts = opts + } else { + if opts.Name != "" { + f.meta.Opts.Name = opts.Name + } + if opts.Path != "" { + f.meta.Opts.Path = opts.Path + } + } } - if urlChanged { - // The playlist URL changed (e.g. refreshed signed link): drop the plan - // so the engine re-resolves on next start (it does when Res is nil). - f.state = nil - f.meta.Res = nil + f.reloadExtraLocked() + return nil +} + +// mergeReqExtra merges the typed extra of incoming into cur following the +// HTTP protocol semantics: Method/Body are overridden when non-empty and +// headers are merged per key. +func mergeReqExtra(cur *base.Request, incoming *base.Request) error { + if incoming.Extra == nil { + return nil } - return + if err := base.ParseReqExtra[ReqExtra](incoming); err != nil { + return err + } + inc, ok := incoming.Extra.(*ReqExtra) + if !ok { + return nil + } + if err := base.ParseReqExtra[ReqExtra](cur); err != nil { + return err + } + target, ok := cur.Extra.(*ReqExtra) + if !ok { + target = &ReqExtra{} + cur.Extra = target + } + if inc.Method != "" { + target.Method = inc.Method + } + if inc.Body != "" { + target.Body = inc.Body + } + if len(inc.Header) > 0 { + if target.Header == nil { + target.Header = make(map[string]string, len(inc.Header)) + } + for key, value := range inc.Header { + target.Header[key] = value + } + } + return nil +} + +// dropPlanLocked discards the resolved plan after a URL change. +func (f *Fetcher) dropPlanLocked() { + f.state = nil + f.meta.Res = nil +} + +// existingOutputNameLocked returns the output file name the task already +// uses, so a re-resolve does not rename an in-progress download. +func (f *Fetcher) existingOutputNameLocked() string { + if f.meta.Opts != nil && f.meta.Opts.Name != "" { + return f.meta.Opts.Name + } + if res := f.meta.Res; res != nil && len(res.Files) > 0 && res.Files[0] != nil { + return res.Files[0].Name + } + return "" } func (f *Fetcher) Start() error { @@ -272,8 +381,30 @@ func (f *Fetcher) Start() error { return errors.New("hls fetcher already running") } if f.state == nil { - f.mu.Unlock() - return errors.New("hls fetcher is not resolved") + if f.meta.Req == nil || f.meta.Res == nil { + f.mu.Unlock() + return errors.New("hls fetcher is not resolved") + } + // Crash window: the engine persists the task (with its resource) + // before the first HLS state checkpoint lands, so a restored task + // can have a resource but no plan. Rebuild the plan instead of + // failing, keeping the existing output name. + if f.client == nil { + f.client = f.buildClient() + } + state, _, err := f.resolvePlanLocked(context.Background(), f.existingOutputNameLocked()) + if err != nil { + f.mu.Unlock() + return fmt.Errorf("re-resolve failed: %w", err) + } + f.state = state + size := planTotalSize(state.Segments) + if res := f.meta.Res; res != nil { + res.Size = size + if len(res.Files) > 0 && res.Files[0] != nil { + res.Files[0].Size = size + } + } } // The client is built here, before any worker exists, so segment // goroutines never race on lazy initialization. diff --git a/internal/protocol/hls/fetcher_test.go b/internal/protocol/hls/fetcher_test.go index 34399ec4..3dd785f0 100644 --- a/internal/protocol/hls/fetcher_test.go +++ b/internal/protocol/hls/fetcher_test.go @@ -4,6 +4,7 @@ import ( "bytes" "crypto/aes" "crypto/cipher" + "encoding/json" "fmt" "net/http" "net/http/httptest" @@ -480,3 +481,377 @@ func hexEncode(b []byte) string { } return string(out) } + +// restoreLikeEngine round-trips the persisted state and the task meta through +// JSON the way the engine storage does, then builds a fresh fetcher. This is +// the path that historically lost the request extra headers. +func restoreLikeEngine(t *testing.T, f *Fetcher, cfg config) *Fetcher { + t.Helper() + fm := &FetcherManager{} + stateAny, err := fm.Store(f) + if err != nil { + t.Fatal(err) + } + rawState, err := json.Marshal(stateAny) + if err != nil { + t.Fatal(err) + } + storedState := &fetcherState{} + if err := json.Unmarshal(rawState, storedState); err != nil { + t.Fatal(err) + } + rawMeta, err := json.Marshal(f.meta) + if err != nil { + t.Fatal(err) + } + restoredMeta := &fetcher.FetcherMeta{} + if err := json.Unmarshal(rawMeta, restoredMeta); err != nil { + t.Fatal(err) + } + if _, ok := restoredMeta.Req.Extra.(*ReqExtra); ok { + t.Fatal("storage round trip should yield a map-shaped extra, like the engine") + } + ctl := controller.NewController() + ctl.GetConfig = func(v any) { + if c, ok := v.(*config); ok { + *c = cfg + } + } + _, builder := fm.Restore() + rf := builder(restoredMeta, storedState).(*Fetcher) + // The engine calls Setup on the builder result (restoreFetcher). + rf.Setup(ctl) + return rf +} + +func TestFetcherE2E_StoreRestoreAuthenticated(t *testing.T) { + const total = 4 + page := "https://player.example.com/watch/42" + var served atomic.Int64 + mux := http.NewServeMux() + var playlist strings.Builder + playlist.WriteString("#EXTM3U\n") + for i := 0; i < total; i++ { + playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i)) + } + playlist.WriteString("#EXT-X-ENDLIST\n") + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(playlist.String())) + }) + segments := make(map[int][]byte) + for i := 0; i < total; i++ { + segments[i] = []byte(fmt.Sprintf("AUTH-SEGMENT-%02d", i)) + idx := i + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", idx), func(w http.ResponseWriter, r *http.Request) { + time.Sleep(80 * time.Millisecond) + served.Add(1) + w.Write(segments[idx]) + }) + } + server := httptest.NewServer(&cdnHandler{mux: mux, requireReferer: page}) + defer server.Close() + + cfg := testConfig() + cfg.SegmentConnections = 2 + cfg.PrefetchContentLength = false + + f := newTestFetcher(t, cfg) + req := &base.Request{ + URL: server.URL + "/video.m3u8", + Extra: &ReqExtra{Header: map[string]string{"Referer": page}}, + } + if err := f.Resolve(req, &base.Options{Path: t.TempDir()}); err != nil { + t.Fatal(err) + } + if err := f.Start(); err != nil { + t.Fatal(err) + } + time.Sleep(150 * time.Millisecond) + if err := f.Pause(); err != nil { + t.Fatal(err) + } + + rf := restoreLikeEngine(t, f, cfg) + if rf.extra == nil || rf.extra.Header["Referer"] != page { + t.Fatalf("restored fetcher lost the request extra headers: %+v", rf.extra) + } + if err := rf.Start(); err != nil { + t.Fatal(err) + } + if err := rf.Wait(); err != nil { + t.Fatalf("resumed download failed (request headers lost?): %v", err) + } + var want []byte + for i := 0; i < total; i++ { + want = append(want, segments[i]...) + } + if !bytes.Equal(readOutput(t, rf), want) { + t.Errorf("restored output mismatch: %q", readOutput(t, rf)) + } +} + +func TestFetcherE2E_PatchPartialUpdate(t *testing.T) { + page := "https://player.example.com/watch/7" + var requiredReferer atomic.Value // string + requiredReferer.Store(page) + var playlist strings.Builder + playlist.WriteString("#EXTM3U\n") + segments := make(map[string][]byte) + for i := 0; i < 3; i++ { + segments[fmt.Sprintf("/s%d.ts", i)] = []byte(fmt.Sprintf("PATCH-%d", i)) + playlist.WriteString(fmt.Sprintf("#EXTINF:4,\ns%d.ts\n", i)) + } + playlist.WriteString("#EXT-X-ENDLIST\n") + mux := http.NewServeMux() + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(playlist.String())) + }) + for path, content := range segments { + c := content + mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + w.Write(c) + }) + } + // The Referer requirement can be flipped mid-test to prove that a patched + // header is actually sent by the next run. + dyn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if want, _ := requiredReferer.Load().(string); want != "" && r.Header.Get("Referer") != want { + w.WriteHeader(http.StatusForbidden) + return + } + mux.ServeHTTP(w, r) + }) + server := httptest.NewServer(dyn) + defer server.Close() + + f := newTestFetcher(t, testConfig()) + req := &base.Request{ + URL: server.URL + "/video.m3u8", + Extra: &ReqExtra{Header: map[string]string{"Referer": page}}, + } + opts := &base.Options{Path: t.TempDir()} + if err := f.Resolve(req, opts); err != nil { + t.Fatal(err) + } + + // A nil/nil patch must be a no-op, not a panic. + if err := f.Patch(nil, nil); err != nil { + t.Fatalf("nil patch should be accepted: %v", err) + } + + // Header-only patch: the resolved plan and output options are retained + // while the new header takes effect on the next run. + newReferer := "https://player.example.com/watch/9" + if err := f.Patch(&base.Request{ + URL: server.URL + "/video.m3u8", + Extra: &ReqExtra{Header: map[string]string{"Referer": newReferer}}, + }, nil); err != nil { + t.Fatal(err) + } + if f.state == nil || f.meta.Res == nil { + t.Fatal("header-only patch must keep the resolved plan") + } + if f.meta.Opts.Path != opts.Path { + t.Fatalf("header-only patch must keep the output path: %s", f.meta.Opts.Path) + } + if f.extra.Header["Referer"] != newReferer { + t.Fatalf("patched header not applied: %+v", f.extra.Header) + } + requiredReferer.Store(newReferer) + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Wait(); err != nil { + t.Fatalf("run with patched header failed: %v", err) + } + if got := readOutput(t, f); string(got) != "PATCH-0PATCH-1PATCH-2" { + t.Errorf("output mismatch after header patch: %q", got) + } + + // Options-only patch: renaming the task keeps the resolved plan. + if err := f.Patch(nil, &base.Options{Name: "renamed.ts"}); err != nil { + t.Fatal(err) + } + if f.state == nil { + t.Fatal("options-only patch must keep the resolved plan") + } + if f.meta.Opts.Name != "renamed.ts" || f.meta.Opts.Path != opts.Path { + t.Fatalf("options patch misapplied: %+v", f.meta.Opts) + } + + // URL-only patch: the plan is dropped for a re-resolve but the output + // path and renamed file survive. + if err := f.Patch(&base.Request{URL: server.URL + "/refreshed.m3u8"}, nil); err != nil { + t.Fatal(err) + } + if f.state != nil || f.meta.Res != nil { + t.Fatal("URL patch must drop the resolved plan") + } + if f.meta.Opts.Path != opts.Path || f.meta.Opts.Name != "renamed.ts" { + t.Fatal("URL patch must keep the output options") + } +} + +func TestFetcherE2E_PatchWhileRunningRejected(t *testing.T) { + mux := http.NewServeMux() + var playlist strings.Builder + playlist.WriteString("#EXTM3U\n") + for i := 0; i < 6; i++ { + playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i)) + } + playlist.WriteString("#EXT-X-ENDLIST\n") + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(playlist.String())) + }) + for i := 0; i < 6; i++ { + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", i), func(w http.ResponseWriter, r *http.Request) { + time.Sleep(150 * time.Millisecond) + w.Write([]byte("X")) + }) + } + server := httptest.NewServer(mux) + defer server.Close() + + f := newTestFetcher(t, testConfig()) + if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil { + t.Fatal(err) + } + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Patch(&base.Request{URL: server.URL + "/other.m3u8"}, nil); err == nil { + t.Fatal("patching a running task must be rejected") + } + if err := f.Pause(); err != nil { + t.Fatal(err) + } + // After the pause the patch goes through and only then drops the plan. + if err := f.Patch(&base.Request{URL: server.URL + "/other.m3u8"}, nil); err != nil { + t.Fatal(err) + } + if f.state != nil { + t.Fatal("URL patch after pause must drop the plan") + } +} + +func TestFetcherE2E_RetryAfterFailure(t *testing.T) { + const total = 4 + var healthy atomic.Bool + mux := http.NewServeMux() + var playlist strings.Builder + playlist.WriteString("#EXTM3U\n") + for i := 0; i < total; i++ { + playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i)) + } + playlist.WriteString("#EXT-X-ENDLIST\n") + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(playlist.String())) + }) + segments := make(map[int][]byte) + for i := 0; i < total; i++ { + segments[i] = []byte(fmt.Sprintf("RETRY-SEGMENT-%02d", i)) + idx := i + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", idx), func(w http.ResponseWriter, r *http.Request) { + if idx == 1 && !healthy.Load() { + w.WriteHeader(http.StatusInternalServerError) + return + } + time.Sleep(60 * time.Millisecond) + w.Write(segments[idx]) + }) + } + server := httptest.NewServer(mux) + defer server.Close() + + cfg := testConfig() + cfg.MaxRetries = 0 // fail the run on the first bad segment + cfg.PrefetchContentLength = false + f := newTestFetcher(t, cfg) + if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil { + t.Fatal(err) + } + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Wait(); err == nil { + t.Fatal("first run should fail while the segment endpoint is down") + } + if got := statsOf(t, f).SegmentsFailed; got == 0 { + t.Fatal("failed segments should be reported in stats") + } + + // Server recovered: a fresh run must succeed and report zero failures. + healthy.Store(true) + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Wait(); err != nil { + t.Fatalf("retry run failed: %v", err) + } + stats := statsOf(t, f) + if stats.SegmentsFailed != 0 { + t.Errorf("failures must reset per run, got %d", stats.SegmentsFailed) + } + if stats.SegmentsDone != total { + t.Errorf("want %d done, got %d", total, stats.SegmentsDone) + } + var want []byte + for i := 0; i < total; i++ { + want = append(want, segments[i]...) + } + if !bytes.Equal(readOutput(t, f), want) { + t.Errorf("retry output mismatch: %q", readOutput(t, f)) + } +} + +func TestFetcherE2E_StateMissingSelfHeal(t *testing.T) { + const total = 3 + mux := http.NewServeMux() + var playlist strings.Builder + playlist.WriteString("#EXTM3U\n") + for i := 0; i < total; i++ { + playlist.WriteString(fmt.Sprintf("#EXTINF:4,\nseg-%d.ts\n", i)) + } + playlist.WriteString("#EXT-X-ENDLIST\n") + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(playlist.String())) + }) + segments := make(map[int][]byte) + for i := 0; i < total; i++ { + segments[i] = []byte(fmt.Sprintf("SELFHEAL-%02d", i)) + idx := i + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", idx), func(w http.ResponseWriter, r *http.Request) { + w.Write(segments[idx]) + }) + } + server := httptest.NewServer(mux) + defer server.Close() + + f := newTestFetcher(t, testConfig()) + if err := f.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: t.TempDir()}); err != nil { + t.Fatal(err) + } + // Simulate the crash window where the engine persisted the task with its + // resource but the HLS plan was never checkpointed. + f.state = nil + outputName := f.meta.Res.Files[0].Name + if err := f.Start(); err != nil { + t.Fatalf("start should self-heal a missing plan: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("self-healed run failed: %v", err) + } + if f.state == nil || len(f.state.Segments) != total { + t.Fatal("self-heal should rebuild the segment plan") + } + if f.meta.Res.Files[0].Name != outputName { + t.Fatalf("self-heal renamed the output: %s", f.meta.Res.Files[0].Name) + } + var want []byte + for i := 0; i < total; i++ { + want = append(want, segments[i]...) + } + if !bytes.Equal(readOutput(t, f), want) { + t.Errorf("self-healed output mismatch: %q", readOutput(t, f)) + } +}