diff --git a/internal/protocol/hls/fetcher.go b/internal/protocol/hls/fetcher.go index 9fc041b9..a178c0db 100644 --- a/internal/protocol/hls/fetcher.go +++ b/internal/protocol/hls/fetcher.go @@ -3,7 +3,6 @@ package hls import ( "bytes" "context" - "encoding/json" "errors" "fmt" "io" @@ -46,11 +45,17 @@ func copyReqExtra(extra *ReqExtra) *ReqExtra { return copied } -// fetcherState is the resolved download plan persisted for resume. +// fetcherState is the resolved download plan persisted for resume. StagingID +// gives the task its own staging folder (two tasks downloading the same URL +// into the same directory never share segment files) and survives restarts +// through the persisted state. StagingDir pins the absolute folder so user +// renames or moves of the task output do not orphan the staged segments. type fetcherState struct { MediaURL string `json:"mediaURL"` OutputName string `json:"outputName"` Segments []*Segment `json:"segments"` + StagingID string `json:"stagingID,omitempty"` + StagingDir string `json:"stagingDir,omitempty"` } // Stats is the protocol task statistics snapshot. @@ -86,6 +91,8 @@ type hlsRun struct { keyCache map[string][]byte keyCacheMu sync.Mutex + stagingID string // identity of the staging folder, bound into the journal + planHash string // fingerprint of the download plan, validated on resume journal map[int64]int64 journalPath string journalMu sync.Mutex @@ -184,7 +191,12 @@ func (f *Fetcher) Resolve(req *base.Request, opts *base.Options) (err error) { if err != nil { return err } + state.StagingID = randomStagingID() f.state = state + // Pin the staging folder now so the very first persisted state already + // carries it: a crash right after create still resumes into the same + // folder. + f.resolveStagingDirLocked() size := planTotalSize(state.Segments) f.meta.Res = &base.Resource{ Range: false, @@ -352,8 +364,13 @@ func mergeReqExtra(cur *base.Request, incoming *base.Request) error { return nil } -// dropPlanLocked discards the resolved plan after a URL change. +// dropPlanLocked discards the resolved plan after a URL change. The staging +// folder of the dropped plan is removed too: this only runs while the +// fetcher is quiescent, so no run can still be writing into it. func (f *Fetcher) dropPlanLocked() { + if f.state != nil && f.state.StagingDir != "" { + os.RemoveAll(f.state.StagingDir) + } f.state = nil f.meta.Res = nil } @@ -421,7 +438,9 @@ func (f *Fetcher) Start() error { extra: copyReqExtra(f.extra), client: f.client, outputPath: f.meta.SingleFilepath(), - tempDir: f.stagingDirLocked(), + stagingID: f.state.StagingID, + planHash: planFingerprint(f.state), + tempDir: f.resolveStagingDirLocked(), keyCache: make(map[string][]byte), journal: make(map[int64]int64), } @@ -464,7 +483,7 @@ func (f *Fetcher) Close() error { run := f.activeRun tempDir := "" if f.state != nil { - tempDir = f.stagingDirLocked() + tempDir = f.resolveStagingDirLocked() } f.mu.Unlock() @@ -757,15 +776,20 @@ func (f *Fetcher) fetchKey(run *hlsRun, key *Key) ([]byte, error) { } // merge concatenates the downloaded segments in plan order into the final -// output file. Segments are stored decrypted, so this is pure I/O. The final -// size is published through Progress: the engine backfills Res.Size from -// Progress().TotalDownloaded() on completion. +// output file. Segments are stored decrypted, so this is pure I/O. The merge +// writes to a unique temporary file next to the target, checks write and +// close errors, and only then replaces the target: a failed merge leaves the +// staged segments intact so it can be retried, and the unique part name keeps +// concurrent tasks that share an output path from clobbering each other. The +// final size is published through Progress: the engine backfills Res.Size +// from Progress().TotalDownloaded() on completion. func (f *Fetcher) merge(run *hlsRun) error { target := run.outputPath if err := os.MkdirAll(filepath.Dir(target), 0777); err != nil { return err } - out, err := os.Create(target) + partPath := target + "." + tempDirPrefix + run.stagingID + ".part" + out, err := os.Create(partPath) if err != nil { return err } @@ -774,15 +798,29 @@ func (f *Fetcher) merge(run *hlsRun) error { data, err := os.ReadFile(segmentPath(run.tempDir, int64(i))) if err != nil { out.Close() + os.Remove(partPath) return fmt.Errorf("read segment %d failed: %w", i, err) } if _, err := out.Write(data); err != nil { out.Close() + os.Remove(partPath) return err } total += int64(len(data)) } if err := out.Close(); err != nil { + os.Remove(partPath) + return err + } + // Remove-then-rename keeps this idempotent: if a previous merge already + // renamed its output before a crash lost the bookkeeping, re-merging + // must still succeed. + if err := os.Remove(target); err != nil && !os.IsNotExist(err) { + os.Remove(partPath) + return err + } + if err := os.Rename(partPath, target); err != nil { + os.Remove(partPath) return err } run.downloaded.Store(total) @@ -807,45 +845,27 @@ func segmentPath(tempDir string, idx int64) string { return filepath.Join(tempDir, "seg-"+strconv.FormatInt(idx, 10)) } -// stagingDirLocked derives the staging folder for the current state. It must -// be called with f.mu held. -func (f *Fetcher) stagingDirLocked() string { - return filepath.Join(filepath.Dir(f.meta.SingleFilepath()), tempDirPrefix+sha1Hex(f.state.MediaURL)[:12]) -} - -func loadJournal(run *hlsRun) error { - data, err := os.ReadFile(run.journalPath) - if err != nil { - if os.IsNotExist(err) { - return nil +// resolveStagingDirLocked returns the staging folder of the current plan, +// deriving and pinning it on first use so later output renames or moves do +// not orphan the staged segments. It lives next to the task output, which is +// exactly opts.Path for a single-file task. Must be called with f.mu held. +func (f *Fetcher) resolveStagingDirLocked() string { + if f.state.StagingDir == "" { + if f.state.StagingID == "" { + // Legacy state (or a hand-crafted plan): migrate to a fresh task + // owned identity. The old URL-hash folder is deliberately not + // reused: its contents cannot be proven to belong to this plan, + // and sharing it with another same-URL task is exactly the bug + // being fixed. It is left behind as orphaned temp data. + f.state.StagingID = randomStagingID() } - return err - } - stored := make(map[string]int64) - if err := json.Unmarshal(data, &stored); err != nil { - // A broken journal must not kill a restartable download. - return nil - } - for key, size := range stored { - idx, err := strconv.ParseInt(key, 10, 64) - if err != nil { - continue + dir := "" + if f.meta.Opts != nil { + dir = f.meta.Opts.Path } - run.journal[idx] = size + f.state.StagingDir = filepath.Join(dir, tempDirPrefix+f.state.StagingID) } - return nil -} - -func saveJournal(run *hlsRun) error { - stored := make(map[string]int64, len(run.journal)) - for idx, size := range run.journal { - stored[strconv.FormatInt(idx, 10)] = size - } - data, err := json.Marshal(stored) - if err != nil { - return err - } - return os.WriteFile(run.journalPath, data, 0666) + return f.state.StagingDir } // prefetchSize probes segment Content-Length via HEAD. Any failure resets all diff --git a/internal/protocol/hls/journal.go b/internal/protocol/hls/journal.go new file mode 100644 index 00000000..3de175a0 --- /dev/null +++ b/internal/protocol/hls/journal.go @@ -0,0 +1,110 @@ +package hls + +import ( + "crypto/rand" + "crypto/sha1" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "strconv" +) + +// journalFile is the persisted completed-segment ledger. It is bound to the +// exact download plan: a journal whose fingerprint does not match the current +// plan is discarded instead of being applied to the wrong segments. +type journalFile struct { + StagingID string `json:"stagingID"` + PlanHash string `json:"planHash"` + SegmentCount int `json:"segmentCount"` + Sizes map[string]int64 `json:"sizes"` +} + +// randomStagingID returns a random identifier used to give every task its own +// staging folder, stable across app restarts through the persisted state. +func randomStagingID() string { + buf := make([]byte, 6) + if _, err := rand.Read(buf); err != nil { + // crypto/rand failing is catastrophic anyway; fall back to address + // entropy rather than crashing the fetch. + return fmt.Sprintf("%x", sha1.Sum([]byte(fmt.Sprintf("%p", buf))))[:12] + } + return hex.EncodeToString(buf) +} + +// planFingerprint hashes everything that defines the download plan: the media +// URL, the ordered segments with their sequence numbers, byte ranges, keys +// and init markers. Two tasks resolve to different fingerprints even when +// they point at the same playlist URL. +func planFingerprint(state *fetcherState) string { + h := sha1.New() + io.WriteString(h, state.MediaURL+"\n") + io.WriteString(h, strconv.Itoa(len(state.Segments))+"\n") + for _, seg := range state.Segments { + fmt.Fprintf(h, "%d|%s|", seg.Sequence, seg.URI) + if seg.Byterange != nil { + fmt.Fprintf(h, "%d@%d|", seg.Byterange.Offset, seg.Byterange.Length) + } + if seg.Key != nil { + fmt.Fprintf(h, "key=%s|", seg.Key.URI) + } + if seg.IsInit { + io.WriteString(h, "init|") + } + io.WriteString(h, "\n") + } + return hex.EncodeToString(h.Sum(nil))[:16] +} + +// loadJournal reads the journal of a run, discarding anything that cannot be +// proven to belong to the current plan: foreign or legacy journals, corrupted +// files, and entries whose recorded size no longer matches the plan. +func loadJournal(run *hlsRun) error { + data, err := os.ReadFile(run.journalPath) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + var file journalFile + if err := json.Unmarshal(data, &file); err != nil { + // A broken or legacy journal must not kill a restartable download. + return nil + } + if file.StagingID != run.stagingID || file.PlanHash != run.planHash || file.SegmentCount != len(run.segs) { + return nil + } + for key, size := range file.Sizes { + idx, err := strconv.ParseInt(key, 10, 64) + if err != nil || idx < 0 || idx >= int64(len(run.segs)) || size <= 0 { + continue + } + run.journal[idx] = size + } + return nil +} + +// saveJournal atomically rewrites the journal through a temp file + rename so +// a crash never leaves a half-written ledger behind. +func saveJournal(run *hlsRun) error { + file := journalFile{ + StagingID: run.stagingID, + PlanHash: run.planHash, + SegmentCount: len(run.segs), + Sizes: make(map[string]int64, len(run.journal)), + } + for idx, size := range run.journal { + file.Sizes[strconv.FormatInt(idx, 10)] = size + } + data, err := json.Marshal(file) + if err != nil { + return err + } + tmp := run.journalPath + ".tmp" + if err := os.WriteFile(tmp, data, 0666); err != nil { + return err + } + return os.Rename(tmp, run.journalPath) +} diff --git a/internal/protocol/hls/staging_test.go b/internal/protocol/hls/staging_test.go new file mode 100644 index 00000000..ec25acbf --- /dev/null +++ b/internal/protocol/hls/staging_test.go @@ -0,0 +1,296 @@ +package hls + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/GopeedLab/gopeed/pkg/base" +) + +func TestFetcherE2E_SameURLTwoTasks(t *testing.T) { + const total = 6 + 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("SHARED-SEGMENT-%02d", i)) + idx := i + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", idx), func(w http.ResponseWriter, r *http.Request) { + time.Sleep(120 * time.Millisecond) + w.Write(segments[idx]) + }) + } + server := httptest.NewServer(mux) + defer server.Close() + + cfg := testConfig() + cfg.SegmentConnections = 2 + cfg.PrefetchContentLength = false + dir := t.TempDir() + + // Two tasks for the same playlist URL into the same directory, like a + // user duplicating a task. The engine would rename the second output + // (AutoRename); mirror that here. + f1 := newTestFetcher(t, cfg) + if err := f1.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: dir}); err != nil { + t.Fatal(err) + } + f2 := newTestFetcher(t, cfg) + if err := f2.Resolve(&base.Request{URL: server.URL + "/video.m3u8"}, &base.Options{Path: dir, Name: "video-2.ts"}); err != nil { + t.Fatal(err) + } + if f1.state.StagingDir == "" || f1.state.StagingDir == f2.state.StagingDir { + t.Fatalf("staging dirs must be distinct per task: %q vs %q", f1.state.StagingDir, f2.state.StagingDir) + } + + // Start task 1, let it stage some segments, then pause it. + if err := f1.Start(); err != nil { + t.Fatal(err) + } + time.Sleep(250 * time.Millisecond) + if err := f1.Pause(); err != nil { + t.Fatal(err) + } + staged1, err := os.ReadDir(f1.state.StagingDir) + if err != nil || len(staged1) == 0 { + t.Fatalf("task 1 should have staged segments, dir: %s", f1.state.StagingDir) + } + + // Task 2 must complete without touching task 1's staging data... + if err := f2.Start(); err != nil { + t.Fatal(err) + } + if err := f2.Wait(); err != nil { + t.Fatal(err) + } + stillThere, err := os.ReadDir(f1.state.StagingDir) + if err != nil || len(stillThere) != len(staged1) { + t.Fatalf("task 2 disturbed task 1 staging: had %d entries, now %d", len(staged1), len(stillThere)) + } + + // ...and task 1 resumes to a complete, correct output of its own. + if err := f1.Start(); err != nil { + t.Fatal(err) + } + if err := f1.Wait(); err != nil { + t.Fatal(err) + } + var want []byte + for i := 0; i < total; i++ { + want = append(want, segments[i]...) + } + if !bytes.Equal(readOutput(t, f1), want) || !bytes.Equal(readOutput(t, f2), want) { + t.Error("same-URL task outputs mismatch") + } + entries, _ := os.ReadDir(dir) + if len(entries) != 2 { + names := make([]string, 0, len(entries)) + for _, e := range entries { + names = append(names, e.Name()) + } + t.Errorf("expected only the two outputs to remain, got %v", names) + } +} + +func TestFetcherE2E_ForeignJournalDiscarded(t *testing.T) { + const total = 3 + 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())) + }) + for i := 0; i < total; i++ { + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", i), func(w http.ResponseWriter, r *http.Request) { + served.Add(1) + w.Write([]byte("DATA")) + }) + } + server := httptest.NewServer(mux) + defer server.Close() + + cfg := testConfig() + 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) + } + // Plant a journal that does not belong to this plan (wrong fingerprint). + tempDir := f.resolveStagingDirLocked() + if err := os.MkdirAll(tempDir, 0777); err != nil { + t.Fatal(err) + } + foreign := journalFile{ + StagingID: "some-other-task", + PlanHash: "deadbeef", + SegmentCount: total, + Sizes: map[string]int64{"0": 4, "1": 4, "2": 4}, + } + data, _ := json.Marshal(foreign) + if err := os.WriteFile(filepath.Join(tempDir, "journal.json"), data, 0666); err != nil { + t.Fatal(err) + } + + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Wait(); err != nil { + t.Fatal(err) + } + if got := served.Load(); got != total { + t.Errorf("foreign journal must be discarded, want %d downloads, got %d", total, got) + } +} + +func TestFetcherE2E_MergeFailurePreservesStaging(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/video.m3u8", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("#EXTM3U\n#EXTINF:4,\na.ts\n#EXTINF:4,\nb.ts\n#EXT-X-ENDLIST\n")) + }) + mux.HandleFunc("/a.ts", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("AAAA")) }) + mux.HandleFunc("/b.ts", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("BBBB")) }) + 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) + } + run := &hlsRun{ + ctx: context.Background(), + cancel: func() {}, + done: make(chan struct{}), + segs: f.state.Segments, + outputPath: f.meta.SingleFilepath(), + stagingID: f.state.StagingID, + planHash: planFingerprint(f.state), + tempDir: f.resolveStagingDirLocked(), + keyCache: make(map[string][]byte), + journal: make(map[int64]int64), + } + run.journalPath = filepath.Join(run.tempDir, "journal.json") + if err := os.MkdirAll(run.tempDir, 0777); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(segmentPath(run.tempDir, 0), []byte("AAAA"), 0666); err != nil { + t.Fatal(err) + } + // Segment 1 is missing on disk: the merge must fail without deleting the + // staged data or leaving a partial output behind. + if err := f.merge(run); err == nil { + t.Fatal("merge with a missing segment should fail") + } + if _, err := os.Stat(segmentPath(run.tempDir, 0)); err != nil { + t.Fatal("merge failure must preserve staged segments") + } + entries, _ := os.ReadDir(filepath.Dir(run.outputPath)) + for _, e := range entries { + if strings.HasSuffix(e.Name(), ".part") { + t.Errorf("failed merge left a partial output behind: %s", e.Name()) + } + } + + // Recover the segment: the retried merge succeeds. + if err := os.WriteFile(segmentPath(run.tempDir, 1), []byte("BBBB"), 0666); err != nil { + t.Fatal(err) + } + if err := f.merge(run); err != nil { + t.Fatalf("retried merge failed: %v", err) + } + got, err := os.ReadFile(run.outputPath) + if err != nil { + t.Fatal(err) + } + if string(got) != "AAAABBBB" { + t.Errorf("retried merge output mismatch: %q", got) + } +} + +func TestFetcherE2E_ReMergeOverExistingOutput(t *testing.T) { + const total = 2 + 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())) + }) + for i := 0; i < total; i++ { + mux.HandleFunc(fmt.Sprintf("/seg-%d.ts", i), func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("DATA")) + }) + } + 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) + } + tempDir := f.resolveStagingDirLocked() + if err := os.MkdirAll(tempDir, 0777); err != nil { + t.Fatal(err) + } + // Simulate the crash window after a completed rename but before the task + // was marked done: a stale output file, a complete journal and all + // segments staged. + stale := filepath.Join(filepath.Dir(f.meta.SingleFilepath()), f.state.OutputName) + if err := os.WriteFile(stale, []byte("STALE-OUTPUT"), 0666); err != nil { + t.Fatal(err) + } + jr := &hlsRun{ + segs: f.state.Segments, + stagingID: f.state.StagingID, + planHash: planFingerprint(f.state), + journalPath: filepath.Join(tempDir, "journal.json"), + journal: map[int64]int64{0: 4, 1: 4}, + } + if err := saveJournal(jr); err != nil { + t.Fatal(err) + } + for i := 0; i < total; i++ { + if err := os.WriteFile(segmentPath(tempDir, int64(i)), []byte("DATA"), 0666); err != nil { + t.Fatal(err) + } + } + + if err := f.Start(); err != nil { + t.Fatal(err) + } + if err := f.Wait(); err != nil { + t.Fatalf("re-merge over an existing output failed: %v", err) + } + got, err := os.ReadFile(stale) + if err != nil { + t.Fatal(err) + } + if string(got) != "DATADATA" { + t.Errorf("stale output was not replaced: %q", got) + } +}