diff --git a/cmd/gopeed/flags.go b/cmd/gopeed/flags.go index 7acaf9dd..d03d90eb 100644 --- a/cmd/gopeed/flags.go +++ b/cmd/gopeed/flags.go @@ -7,9 +7,11 @@ import ( ) type args struct { - url string - connections *int - dir *string + url string + connections *int + dir *string + checksum *string + checksumAlgo *string } func parse() *args { @@ -20,6 +22,8 @@ func parse() *args { var args args args.connections = flag.Int("C", 16, "Concurrent connections.") args.dir = flag.String("D", dir, "Store directory.") + args.checksum = flag.String("checksum", "", "Expected checksum hash (hex encoded).") + args.checksumAlgo = flag.String("checksum-algo", "md5", "Checksum algorithm: md5, sha1, or sha256.") flag.Parse() t := flag.Args() if len(t) > 0 { diff --git a/cmd/gopeed/main.go b/cmd/gopeed/main.go index cd4e0587..c6b96705 100644 --- a/cmd/gopeed/main.go +++ b/cmd/gopeed/main.go @@ -19,6 +19,18 @@ func main() { var wg sync.WaitGroup wg.Add(1) + var checksumOpt *http.ChecksumOption + if args.checksum != nil && *args.checksum != "" { + algo := "md5" + if args.checksumAlgo != nil && *args.checksumAlgo != "" { + algo = *args.checksumAlgo + } + checksumOpt = &http.ChecksumOption{ + Algorithm: algo, + Expected: *args.checksum, + } + } + _, err := download.Boot(). URL(args.url). Listener(func(event *download.Event) { @@ -43,8 +55,11 @@ func main() { } }). Create(&base.Options{ - Path: *args.dir, - Extra: http.OptsExtra{Connections: *args.connections}, + Path: *args.dir, + Extra: http.OptsExtra{ + Connections: *args.connections, + Checksum: checksumOpt, + }, }) if err != nil { panic(err) diff --git a/internal/protocol/http/fetcher.go b/internal/protocol/http/fetcher.go index 9eca0089..1297ada4 100644 --- a/internal/protocol/http/fetcher.go +++ b/internal/protocol/http/fetcher.go @@ -2,8 +2,13 @@ package http import ( "context" + "crypto/md5" + "crypto/sha1" + "crypto/sha256" + "encoding/hex" "errors" "fmt" + "hash" "io" "net/http" "net/url" @@ -280,6 +285,12 @@ type Fetcher struct { primaryReadyOnce sync.Once primaryReadyCh chan struct{} + // checksumFailed marks that the last completed download failed + // checksum verification, so the next retry must force a full + // re-download instead of skipping connections that already + // reported Completed. + checksumFailed atomic.Bool + // Start pending mechanism startPending atomic.Bool resolvedCh chan struct{} // Signal when resolve completes @@ -735,12 +746,26 @@ func (f *Fetcher) doStart() error { // If retrying after error, reset connection states for retry if state == stateError { + forceFullReset := f.checksumFailed.Swap(false) + if forceFullReset { + // Checksum mismatch: the file content itself is wrong even + // though every connection reported Completed. Truncate the + // corrupt file so stale bytes cannot linger if the retry is + // interrupted before every connection re-writes its chunk. + if err := os.Truncate(f.meta.SingleFilepath(), 0); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to reset file after checksum mismatch: %w", err) + } + } f.connMu.Lock() for _, conn := range f.connections { - // Reset connections that can be retried - if !conn.Completed && conn.State != connCompleted { + // On a normal error, only reset connections that didn't + // finish. On a checksum mismatch, reset EVERY connection — + // "Completed" only means the download loop finished, not + // that the bytes were verified correct. + shouldReset := forceFullReset || (!conn.Completed && conn.State != connCompleted) + if shouldReset { if !f.hasSequentialPrefixLocked(conn) { - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, forceFullReset) } conn.State = connNotStarted conn.failed = false @@ -970,8 +995,14 @@ func (f *Fetcher) expandConnections() { } f.fileMu.Unlock() - f.setState(stateDone) - f.doneCh <- nil + var err error + if chkErr := f.verifyChecksum(); chkErr != nil { + err = chkErr + f.setState(stateError) + } else { + f.setState(stateDone) + } + f.doneCh <- err return } @@ -1107,7 +1138,7 @@ func (f *Fetcher) runConnection(conn *connection) { // it only when no safe validator exists or If-Range returns a full 200. f.connMu.Lock() if !f.hasSequentialPrefixLocked(conn) { - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, false) } f.connMu.Unlock() } @@ -1214,7 +1245,7 @@ func (f *Fetcher) downloadChunkOnce(conn *connection, client *http.Client, buf [ resumeProbe := f.canProbeSequentialResumeLocked(conn) intentionalRestart := sequentialSizeUnknown if sequentialSizeUnknown || (hasSequentialPrefix && !resumeProbe) { - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, false) f.resolveDataPos.Store(0) f.rangeValidatorPinned = false intentionalRestart = true @@ -1659,7 +1690,7 @@ func (f *Fetcher) fallbackToSequentialDownload(conn *connection, ifRange string) f.rangeReprobeEligible = true f.rangeValidatorPinned = false f.ifRange = ifRange - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, false) f.resolveDataPos.Store(0) return nil } @@ -1675,7 +1706,7 @@ func (f *Fetcher) restartSequentialDownload(conn *connection, ifRange string) { f.rangeReprobeEligible = true f.rangeValidatorPinned = false f.ifRange = ifRange - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, false) f.resolveDataPos.Store(0) } @@ -2058,16 +2089,25 @@ func (f *Fetcher) helpOtherConnection(helper *connection) bool { return true } -func (f *Fetcher) resetConnectionForRestart(conn *connection) { - if f.meta.Res.Range { +func (f *Fetcher) resetConnectionForRestart(conn *connection, force bool) { + if f.meta.Res.Range && !force { return } + if f.meta.Res.Range && force { + // Checksum mismatch: the on-disk content for this chunk is + // confirmed wrong, so its bytes must be re-fetched from the + // origin server even though range support would normally let + // us just resume from the existing Downloaded position. + if conn.Chunk != nil { + conn.Chunk.Downloaded = 0 + } + } // Without range support a new request always starts from byte 0, // so pause/retry must restart instead of continuing from the old offset. if conn.Chunk == nil { conn.Chunk = newChunk(0, 0) - } else { + } else if !f.meta.Res.Range { conn.Chunk.Begin = 0 conn.Chunk.setEnd(0) conn.Chunk.Downloaded = 0 @@ -2131,7 +2171,7 @@ func (f *Fetcher) resumeConnections() { } } if !f.hasSequentialPrefixLocked(conn) { - f.resetConnectionForRestart(conn) + f.resetConnectionForRestart(conn, false) } // Reset the connection state for resume conn.ctx, conn.cancel = context.WithCancel(f.ctx) @@ -2238,10 +2278,15 @@ func (f *Fetcher) onDownloadComplete() { } f.fileMu.Unlock() - if finalErr != nil { - f.setState(stateError) + if finalErr == nil { + if chkErr := f.verifyChecksum(); chkErr != nil { + finalErr = chkErr + f.setState(stateError) + } else { + f.setState(stateDone) + } } else { - f.setState(stateDone) + f.setState(stateError) } select { @@ -2250,6 +2295,51 @@ func (f *Fetcher) onDownloadComplete() { } } +func (f *Fetcher) verifyChecksum() error { + if f.meta == nil || f.meta.Opts == nil || f.meta.Opts.Extra == nil { + return nil + } + extra, ok := f.meta.Opts.Extra.(*fhttp.OptsExtra) + if !ok || extra.Checksum == nil { + return nil + } + chk := extra.Checksum + if chk.Algorithm == "" && chk.Expected == "" { + return nil + } + var h hash.Hash + switch strings.ToLower(strings.TrimSpace(chk.Algorithm)) { + case "md5": + h = md5.New() + case "sha1", "sha-1": + h = sha1.New() + case "sha256", "sha-256": + h = sha256.New() + default: + return fmt.Errorf("unsupported checksum algorithm: %s", chk.Algorithm) + } + + filePath := f.meta.SingleFilepath() + file, err := os.Open(filePath) + if err != nil { + return fmt.Errorf("open file for checksum verification failed: %w", err) + } + defer file.Close() + + buf := make([]byte, 64*1024) + if _, err := io.CopyBuffer(h, file, buf); err != nil { + return fmt.Errorf("calculate checksum failed: %w", err) + } + + actual := hex.EncodeToString(h.Sum(nil)) + expected := strings.ToLower(strings.TrimSpace(chk.Expected)) + if !strings.EqualFold(actual, expected) { + f.checksumFailed.Store(true) + return fmt.Errorf("checksum mismatch: expected %s, got %s", expected, actual) + } + return nil +} + func (f *Fetcher) checkCompletion() bool { // Check if all data has been downloaded f.connMu.Lock() diff --git a/internal/protocol/http/fetcher_checksum_retry_test.go b/internal/protocol/http/fetcher_checksum_retry_test.go new file mode 100644 index 00000000..a7d3573a --- /dev/null +++ b/internal/protocol/http/fetcher_checksum_retry_test.go @@ -0,0 +1,108 @@ +package http + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + gohttp "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + + "github.com/GopeedLab/gopeed/pkg/base" + "github.com/GopeedLab/gopeed/pkg/protocol/http" +) + +func TestFetcher_ChecksumMismatch_Retry(t *testing.T) { + corruptContent := []byte("Corrupt Initial Download Data!") + correctContent := []byte("Correct Retried Download Data!") + + correctSha256Hasher := sha256.New() + correctSha256Hasher.Write(correctContent) + correctSha256 := hex.EncodeToString(correctSha256Hasher.Sum(nil)) + + var reqCount atomic.Int32 + server := httptest.NewServer(gohttp.HandlerFunc(func(w gohttp.ResponseWriter, r *gohttp.Request) { + count := reqCount.Add(1) + var served []byte + if count == 1 { + // First request: serve corrupted content + served = corruptContent + } else { + // Second / subsequent request: serve correct content + served = correctContent + } + w.Header().Set(base.HttpHeaderContentLength, fmt.Sprintf("%d", len(served))) + w.Header().Set(base.HttpHeaderAcceptRanges, base.HttpHeaderBytes) + w.WriteHeader(gohttp.StatusOK) + _, _ = w.Write(served) + })) + defer server.Close() + + downloadDir := t.TempDir() + fileName := "test_checksum_retry.txt" + + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: downloadDir, + Name: fileName, + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "sha256", + Expected: correctSha256, + }, + }, + } + + // 1. Initial attempt -> will download corruptContent, failing checksum verification + if err := f.Resolve(&base.Request{URL: server.URL + "/retry_test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + err := f.Wait() + if err == nil { + t.Fatal("expected checksum mismatch error on initial attempt, got nil") + } + if !strings.Contains(err.Error(), "checksum mismatch") { + t.Fatalf("expected 'checksum mismatch' in error, got %v", err) + } + if f.getState() != stateError { + t.Fatalf("expected stateError after checksum mismatch, got %v", f.getState()) + } + + // 2. Retry download -> should trigger forced full reset, truncate file, and re-download correctContent + if err := f.Start(); err != nil { + t.Fatalf("Start error on retry: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error on retry (expected success): %v", err) + } + if f.getState() != stateDone { + t.Fatalf("expected stateDone after successful retry, got %v", f.getState()) + } + + // 3. Verify on-disk file content matches the re-downloaded bytes + filePath := f.meta.SingleFilepath() + fileBytes, err := os.ReadFile(filePath) + if err != nil { + t.Fatalf("failed to read completed file: %v", err) + } + if string(fileBytes) != string(correctContent) { + t.Fatalf("on-disk file content mismatch:\ngot: %q\nwant: %q", string(fileBytes), string(correctContent)) + } + + // 4. Compute hash independently and verify + actualHasher := sha256.New() + actualHasher.Write(fileBytes) + actualHash := hex.EncodeToString(actualHasher.Sum(nil)) + if actualHash != correctSha256 { + t.Fatalf("re-computed on-disk hash mismatch: got %s, want %s", actualHash, correctSha256) + } +} diff --git a/internal/protocol/http/fetcher_checksum_test.go b/internal/protocol/http/fetcher_checksum_test.go new file mode 100644 index 00000000..3d760011 --- /dev/null +++ b/internal/protocol/http/fetcher_checksum_test.go @@ -0,0 +1,307 @@ +package http + +import ( + "crypto/md5" + "crypto/sha256" + "encoding/hex" + "fmt" + gohttp "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/GopeedLab/gopeed/internal/fetcher" + "github.com/GopeedLab/gopeed/internal/test" + "github.com/GopeedLab/gopeed/pkg/base" + "github.com/GopeedLab/gopeed/pkg/protocol/http" +) + +func TestFetcher_ChecksumVerification(t *testing.T) { + content := []byte("Hello, Gopeed Checksum Verification Test Data!") + md5Hasher := md5.New() + md5Hasher.Write(content) + correctMd5 := hex.EncodeToString(md5Hasher.Sum(nil)) + + sha256Hasher := sha256.New() + sha256Hasher.Write(content) + correctSha256 := hex.EncodeToString(sha256Hasher.Sum(nil)) + + wrongHash := "0123456789abcdef0123456789abcdef" + + server := httptest.NewServer(gohttp.HandlerFunc(func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set(base.HttpHeaderContentLength, fmt.Sprintf("%d", len(content))) + w.Header().Set(base.HttpHeaderAcceptRanges, base.HttpHeaderBytes) + w.WriteHeader(gohttp.StatusOK) + _, _ = w.Write(content) + })) + defer server.Close() + + t.Run("Valid MD5", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_valid_md5.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "md5", + Expected: correctMd5, + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error (expected success): %v", err) + } + }) + + t.Run("Valid SHA256", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_valid_sha256.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "sha256", + Expected: correctSha256, + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error (expected success): %v", err) + } + }) + + t.Run("Wrong MD5 mismatch", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_wrong_md5.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "md5", + Expected: wrongHash, + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + err := f.Wait() + if err == nil { + t.Fatal("Wait expected checksum mismatch error, got nil") + } + if !strings.Contains(err.Error(), "checksum mismatch") { + t.Fatalf("expected 'checksum mismatch' in error, got %v", err) + } + }) + + t.Run("Unsupported Algorithm", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_unsupported.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "invalid-algo", + Expected: wrongHash, + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + err := f.Wait() + if err == nil { + t.Fatal("Wait expected error for unsupported algorithm, got nil") + } + if !strings.Contains(err.Error(), "unsupported checksum algorithm") { + t.Fatalf("expected 'unsupported checksum algorithm' in error, got %v", err) + } + if f.checksumFailed.Load() { + t.Fatal("checksumFailed flag should NOT be set on unsupported algorithm error") + } + }) + + t.Run("Case-Insensitive Hash Comparison", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_case_insensitive.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "sha256", + Expected: strings.ToUpper(correctSha256), + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error (expected success with upper case hash): %v", err) + } + }) + + t.Run("Empty Expected Checksum (Skipped)", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_empty_expected.txt", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "", + Expected: "", + }, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error (expected skip and success): %v", err) + } + }) + + t.Run("Large File Multi-Chunk Streaming (>64KB)", func(t *testing.T) { + largeData := make([]byte, 128*1024) + for i := range largeData { + largeData[i] = byte(i % 256) + } + largeSha256Hasher := sha256.New() + largeSha256Hasher.Write(largeData) + largeSha256 := hex.EncodeToString(largeSha256Hasher.Sum(nil)) + + largeServer := httptest.NewServer(gohttp.HandlerFunc(func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set(base.HttpHeaderContentLength, fmt.Sprintf("%d", len(largeData))) + w.Header().Set(base.HttpHeaderAcceptRanges, base.HttpHeaderBytes) + w.WriteHeader(gohttp.StatusOK) + _, _ = w.Write(largeData) + })) + defer largeServer.Close() + + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: "test_large_file.bin", + Extra: &http.OptsExtra{ + Connections: 1, + Checksum: &http.ChecksumOption{ + Algorithm: "sha256", + Expected: largeSha256, + }, + }, + } + + if err := f.Resolve(&base.Request{URL: largeServer.URL + "/large.bin"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error (expected large file checksum success): %v", err) + } + }) + + t.Run("No Checksum (Opt-in backward compatibility)", func(t *testing.T) { + f := buildFetcher() + defer f.Close() + + opts := &base.Options{ + Path: t.TempDir(), + Name: test.DownloadName, + Extra: &http.OptsExtra{ + Connections: 1, + }, + } + + if err := f.Resolve(&base.Request{URL: server.URL + "/test.txt"}, opts); err != nil { + t.Fatalf("Resolve error: %v", err) + } + if err := f.Start(); err != nil { + t.Fatalf("Start error: %v", err) + } + if err := f.Wait(); err != nil { + t.Fatalf("Wait error: %v", err) + } + }) +} + +func TestResetConnectionForRestart_ForceOnRangeDownload(t *testing.T) { + f := buildFetcher() + f.meta = &fetcher.FetcherMeta{ + Res: &base.Resource{ + Range: true, + }, + } + conn := &connection{ + Chunk: newChunk(0, 1024), + Downloaded: 512, + Completed: true, + } + conn.Chunk.Downloaded = 512 + + // Non-forced reset on range-supported download does not reset Chunk.Downloaded + f.resetConnectionForRestart(conn, false) + if conn.Chunk.Downloaded != 512 { + t.Fatalf("expected Chunk.Downloaded to remain 512 when force is false, got %d", conn.Chunk.Downloaded) + } + + // Forced reset on range-supported download resets Chunk.Downloaded to 0 + f.resetConnectionForRestart(conn, true) + if conn.Chunk.Downloaded != 0 { + t.Fatalf("expected Chunk.Downloaded to be reset to 0 when force is true, got %d", conn.Chunk.Downloaded) + } + if conn.Completed { + t.Fatal("expected Completed to be false after reset") + } +} diff --git a/pkg/download/downloader.go b/pkg/download/downloader.go index 4d60c7ee..89323144 100644 --- a/pkg/download/downloader.go +++ b/pkg/download/downloader.go @@ -1912,6 +1912,7 @@ func (d *Downloader) doPause(task *Task) (err error) { } if handled { d.emit(EventKeyPause, task) + d.triggerWebhooks(WebhookEventDownloadPause, task, nil) } return nil } diff --git a/pkg/download/model.go b/pkg/download/model.go index bf479a0b..85fa4827 100644 --- a/pkg/download/model.go +++ b/pkg/download/model.go @@ -23,6 +23,9 @@ type ResolveResult struct { Res *base.Resource `json:"res"` } +// Task represents a download unit. +// Options flow path: +// REST Request (model.CreateTask/ResolveTask) -> pkg/download Downloader -> Task.Meta.Opts (base.Options) -> Fetcher.Resolve/Start (e.g. internal/protocol/http/fetcher.go) type Task struct { ID string `json:"id"` Protocol string `json:"protocol"` diff --git a/pkg/download/webhook.go b/pkg/download/webhook.go index d223d213..d0facb05 100644 --- a/pkg/download/webhook.go +++ b/pkg/download/webhook.go @@ -21,6 +21,7 @@ type WebhookEvent string const ( WebhookEventDownloadDone WebhookEvent = "DOWNLOAD_DONE" WebhookEventDownloadError WebhookEvent = "DOWNLOAD_ERROR" + WebhookEventDownloadPause WebhookEvent = "DOWNLOAD_PAUSE" ) // WebhookData is the data sent to webhook URLs diff --git a/pkg/download/webhook_test.go b/pkg/download/webhook_test.go index b2b28093..f87f678f 100644 --- a/pkg/download/webhook_test.go +++ b/pkg/download/webhook_test.go @@ -814,6 +814,49 @@ func TestWebhook_TestWebhookUrl_VerifyTestPayload(t *testing.T) { }) } +func TestWebhook_TriggerOnPause(t *testing.T) { + receivedData := make(chan *WebhookData, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var data WebhookData + if err := json.NewDecoder(r.Body).Decode(&data); err != nil { + t.Errorf("Failed to decode data: %v", err) + return + } + receivedData <- &data + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + setupWebhookTest(t, func(downloader *Downloader) { + cfg, _ := downloader.GetConfig() + cfg.Webhook = &base.WebhookConfig{ + Enable: true, + URLs: []string{server.URL}, + } + downloader.PutConfig(cfg) + + task := NewTask() + task.Protocol = "http" + task.Meta = &mockFetcherMeta + + downloader.triggerWebhooks(WebhookEventDownloadPause, task, nil) + + select { + case data := <-receivedData: + if data.Event != WebhookEventDownloadPause { + t.Errorf("Expected event 'DOWNLOAD_PAUSE', got '%s'", data.Event) + } + if data.Payload == nil || data.Payload.Task == nil { + t.Error("Expected payload.task to be present") + } else if data.Payload.Task.ID != task.ID { + t.Errorf("Expected task ID '%s', got '%s'", task.ID, data.Payload.Task.ID) + } + case <-time.After(2 * time.Second): + t.Error("Timeout waiting for webhook") + } + }) +} + func setupWebhookTest(t *testing.T, fn func(downloader *Downloader)) { defaultDownloader.Setup() defaultDownloader.cfg.StorageDir = ".test_storage" diff --git a/pkg/mcpserver/server.go b/pkg/mcpserver/server.go index 1c6a954e..10c6ec53 100644 --- a/pkg/mcpserver/server.go +++ b/pkg/mcpserver/server.go @@ -12,6 +12,7 @@ import ( "github.com/GopeedLab/gopeed/pkg/base" "github.com/GopeedLab/gopeed/pkg/download" + phttp "github.com/GopeedLab/gopeed/pkg/protocol/http" "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -65,12 +66,13 @@ type btRequestExtra struct { } type httpTaskExtra struct { - Connections int `json:"connections,omitempty" jsonschema:"Number of HTTP download connections"` - AutoTorrent *bool `json:"autoTorrent,omitempty" jsonschema:"Automatically create a BitTorrent task after downloading a torrent file"` - DeleteTorrentAfterDownload *bool `json:"deleteTorrentAfterDownload,omitempty" jsonschema:"Delete the torrent file after automatically creating its BitTorrent task"` - AutoExtract *bool `json:"autoExtract,omitempty" jsonschema:"Automatically extract the downloaded archive"` - ArchivePassword string `json:"archivePassword,omitempty" jsonschema:"Password for an encrypted archive"` - DeleteAfterExtract bool `json:"deleteAfterExtract,omitempty" jsonschema:"Delete archive files after successful extraction"` + Connections int `json:"connections,omitempty" jsonschema:"Number of HTTP download connections"` + AutoTorrent *bool `json:"autoTorrent,omitempty" jsonschema:"Automatically create a BitTorrent task after downloading a torrent file"` + DeleteTorrentAfterDownload *bool `json:"deleteTorrentAfterDownload,omitempty" jsonschema:"Delete the torrent file after automatically creating its BitTorrent task"` + AutoExtract *bool `json:"autoExtract,omitempty" jsonschema:"Automatically extract the downloaded archive"` + ArchivePassword string `json:"archivePassword,omitempty" jsonschema:"Password for an encrypted archive"` + DeleteAfterExtract bool `json:"deleteAfterExtract,omitempty" jsonschema:"Delete archive files after successful extraction"` + Checksum *checksumOption `json:"checksum,omitempty" jsonschema:"Optional checksum verification settings"` } // downloadRequest is the MCP-facing form of base.Request. The REST model uses @@ -129,6 +131,11 @@ func isBitTorrentRequest(rawURL string) bool { strings.HasPrefix(upperURL, "DATA:APPLICATION/X-BITTORRENT;BASE64,") } +type checksumOption struct { + Algorithm string `json:"algorithm" jsonschema:"Checksum algorithm: md5, sha1, or sha256"` + Expected string `json:"expected" jsonschema:"Expected hex-encoded checksum hash"` +} + // downloadOptions is the MCP-facing form of base.Options. Only HTTP currently // defines protocol-specific task options, so Extra has one concrete schema. type downloadOptions struct { @@ -186,12 +193,31 @@ func (o *downloadOptions) baseOptions() *base.Options { if o == nil { return nil } + var extra any + if o.Extra != nil { + var checksum *phttp.ChecksumOption + if o.Extra.Checksum != nil { + checksum = &phttp.ChecksumOption{ + Algorithm: o.Extra.Checksum.Algorithm, + Expected: o.Extra.Checksum.Expected, + } + } + extra = &phttp.OptsExtra{ + Connections: o.Extra.Connections, + AutoTorrent: o.Extra.AutoTorrent, + DeleteTorrentAfterDownload: o.Extra.DeleteTorrentAfterDownload, + AutoExtract: o.Extra.AutoExtract, + ArchivePassword: o.Extra.ArchivePassword, + DeleteAfterExtract: o.Extra.DeleteAfterExtract, + Checksum: checksum, + } + } return &base.Options{ Name: o.Name, Path: o.Path, AsDefaultPath: o.AsDefaultPath, SelectFiles: o.SelectFiles, - Extra: o.Extra, + Extra: extra, } } @@ -270,6 +296,13 @@ type getTaskStatsOutput struct { Stats any `json:"stats"` } +type getDownloadSpeedInput struct{} + +type getDownloadSpeedOutput struct { + TotalSpeed int64 `json:"totalSpeed"` + ActiveTaskCount int `json:"activeTaskCount"` +} + type taskActionOutput struct { ID string `json:"id"` Success bool `json:"success"` @@ -416,6 +449,33 @@ func registerTools(server *mcp.Server, downloader *download.Downloader) { return nil, &getTaskStatsOutput{Stats: stats}, nil }) + mcp.AddTool(server, &mcp.Tool{ + Name: "get_download_speed", + Title: "Get download speed", + Description: "Returns the combined download speed across all currently active tasks and how many tasks are actively downloading. Use get_task_status instead if you need the speed, progress, or ETA of one specific task.", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld}, + }, func(_ context.Context, _ *mcp.CallToolRequest, _ *getDownloadSpeedInput) (*mcp.CallToolResult, *getDownloadSpeedOutput, error) { + allTasks := downloader.GetTasks() + var totalSpeed int64 + activeCount := 0 + + for _, task := range allTasks { + status, err := downloader.RuntimeStatus(task.ID) + if err != nil { + continue + } + if status.Status == base.DownloadStatusRunning { + activeCount++ + totalSpeed += status.Speed + } + } + + return nil, &getDownloadSpeedOutput{ + TotalSpeed: totalSpeed, + ActiveTaskCount: activeCount, + }, nil + }) + registerTaskAction(server, downloader, "pause_task", "Pause task", "Pause one Gopeed task.", &nondestructive, &closedWorld, downloader.Pause) registerTaskAction(server, downloader, "continue_task", "Continue task", "Continue a paused or failed Gopeed task.", &nondestructive, &openWorld, downloader.Continue) diff --git a/pkg/mcpserver/server_test.go b/pkg/mcpserver/server_test.go index 56d208e1..225d6ed8 100644 --- a/pkg/mcpserver/server_test.go +++ b/pkg/mcpserver/server_test.go @@ -64,6 +64,7 @@ func TestStreamableHTTPTools(t *testing.T) { "continue_task", "create_task", "delete_task", + "get_download_speed", "get_task", "get_task_stats", "get_task_status", diff --git a/pkg/mcpserver/speed_test.go b/pkg/mcpserver/speed_test.go new file mode 100644 index 00000000..28df1eed --- /dev/null +++ b/pkg/mcpserver/speed_test.go @@ -0,0 +1,103 @@ +package mcpserver + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/GopeedLab/gopeed/pkg/base" + "github.com/GopeedLab/gopeed/pkg/download" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func TestGetDownloadSpeedTool(t *testing.T) { + // Serve a test payload that takes a little bit of time or streams bytes + testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "1048576") + w.WriteHeader(http.StatusOK) + chunk := make([]byte, 8192) + for i := 0; i < 128; i++ { + _, _ = w.Write(chunk) + time.Sleep(10 * time.Millisecond) + } + })) + defer testServer.Close() + + downloader := download.NewDownloader(&download.DownloaderConfig{ + Storage: download.NewMemStorage(), + StorageDir: t.TempDir(), + }) + if err := downloader.Setup(); err != nil { + t.Fatal(err) + } + defer downloader.Close() + + httpServer := httptest.NewServer(NewHandler(downloader)) + defer httpServer.Close() + + client := mcp.NewClient(&mcp.Implementation{Name: "gopeed-test", Version: "1.0.0"}, nil) + session, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatal(err) + } + defer session.Close() + + // 1. Zero tasks test (returns 0 totalSpeed and 0 activeTaskCount) + result, err := session.CallTool(context.Background(), &mcp.CallToolParams{ + Name: "get_download_speed", + Arguments: map[string]any{}, + }) + if err != nil { + t.Fatal(err) + } + if result.IsError { + t.Fatalf("get_download_speed returned error: %+v", result.Content) + } + structured, ok := result.StructuredContent.(map[string]any) + if !ok { + t.Fatalf("structured content type = %T", result.StructuredContent) + } + if activeCount, ok := structured["activeTaskCount"].(float64); !ok || activeCount != 0 { + t.Fatalf("activeTaskCount = %v, want 0", structured["activeTaskCount"]) + } + if totalSpeed, ok := structured["totalSpeed"].(float64); !ok || totalSpeed != 0 { + t.Fatalf("totalSpeed = %v, want 0", structured["totalSpeed"]) + } + + // 2. Start a real download task + taskID, err := downloader.CreateDirect(&base.Request{ + URL: testServer.URL, + }, &base.Options{ + Path: t.TempDir(), + }) + if err != nil { + t.Fatal(err) + } + defer func() { + _ = downloader.Delete(&download.TaskFilter{IDs: []string{taskID}}, true) + }() + + // Query get_download_speed while downloading + result, err = session.CallTool(context.Background(), &mcp.CallToolParams{ + Name: "get_download_speed", + Arguments: map[string]any{}, + }) + if err != nil { + t.Fatal(err) + } + if result.IsError { + t.Fatalf("get_download_speed returned error: %+v", result.Content) + } + structured, ok = result.StructuredContent.(map[string]any) + if !ok { + t.Fatalf("structured content type = %T", result.StructuredContent) + } + if _, ok := structured["activeTaskCount"].(float64); !ok { + t.Fatalf("activeTaskCount missing or wrong type: %v", structured["activeTaskCount"]) + } + if _, ok := structured["totalSpeed"].(float64); !ok { + t.Fatalf("totalSpeed missing or wrong type: %v", structured["totalSpeed"]) + } +} diff --git a/pkg/protocol/http/model.go b/pkg/protocol/http/model.go index 9e96ffe2..c99edf01 100644 --- a/pkg/protocol/http/model.go +++ b/pkg/protocol/http/model.go @@ -6,6 +6,13 @@ type ReqExtra struct { Body string `json:"body"` } +// ChecksumOption specifies an expected checksum to verify against +// the completed download. Only supported for HTTP downloads. +type ChecksumOption struct { + Algorithm string `json:"algorithm"` // "md5" | "sha1" | "sha256" + Expected string `json:"expected"` // hex-encoded expected hash +} + type OptsExtra struct { Connections int `json:"connections"` // AutoTorrent when task download complete, and it is a .torrent file, it will be auto create a new task for the torrent file @@ -21,6 +28,8 @@ type OptsExtra struct { ArchivePassword string `json:"archivePassword"` // DeleteAfterExtract when true, deletes the archive file after successful extraction DeleteAfterExtract bool `json:"deleteAfterExtract"` + // Checksum verifies the completed file against an expected hash after download + Checksum *ChecksumOption `json:"checksum,omitempty"` } // Stats for download diff --git a/pkg/protocol/http/model_test.go b/pkg/protocol/http/model_test.go new file mode 100644 index 00000000..e9348a88 --- /dev/null +++ b/pkg/protocol/http/model_test.go @@ -0,0 +1,15 @@ +package http + +import "testing" + +func TestOptsExtra_Checksum(t *testing.T) { + extra := &OptsExtra{ + Checksum: &ChecksumOption{ + Algorithm: "sha256", + Expected: "abcdef1234567890", + }, + } + if extra.Checksum.Algorithm != "sha256" || extra.Checksum.Expected != "abcdef1234567890" { + t.Fatalf("unexpected checksum: %+v", extra.Checksum) + } +} diff --git a/ui/flutter/lib/api/model/options.dart b/ui/flutter/lib/api/model/options.dart index 17aadd0d..35a50ac2 100644 --- a/ui/flutter/lib/api/model/options.dart +++ b/ui/flutter/lib/api/model/options.dart @@ -17,6 +17,21 @@ class Options { Map toJson() => _$OptionsToJson(this); } +@JsonSerializable() +class ChecksumOption { + String algorithm; + String expected; + + ChecksumOption({ + this.algorithm = '', + this.expected = '', + }); + + factory ChecksumOption.fromJson(Map json) => _$ChecksumOptionFromJson(json); + + Map toJson() => _$ChecksumOptionToJson(this); +} + @JsonSerializable() class OptsExtraHttp { int connections; @@ -25,6 +40,7 @@ class OptsExtraHttp { bool? autoExtract; String archivePassword; bool deleteAfterExtract; + ChecksumOption? checksum; OptsExtraHttp({ this.connections = 0, @@ -33,6 +49,7 @@ class OptsExtraHttp { this.autoExtract, this.archivePassword = '', this.deleteAfterExtract = false, + this.checksum, }); factory OptsExtraHttp.fromJson(Map json) => _$OptsExtraHttpFromJson(json); diff --git a/ui/flutter/lib/api/model/options.g.dart b/ui/flutter/lib/api/model/options.g.dart index 52b86453..557ea04f 100644 --- a/ui/flutter/lib/api/model/options.g.dart +++ b/ui/flutter/lib/api/model/options.g.dart @@ -22,6 +22,16 @@ Map _$OptionsToJson(Options instance) => { 'extra': ?instance.extra, }; +ChecksumOption _$ChecksumOptionFromJson(Map json) => ChecksumOption( + algorithm: json['algorithm'] as String? ?? '', + expected: json['expected'] as String? ?? '', +); + +Map _$ChecksumOptionToJson(ChecksumOption instance) => { + 'algorithm': instance.algorithm, + 'expected': instance.expected, +}; + OptsExtraHttp _$OptsExtraHttpFromJson(Map json) => OptsExtraHttp( connections: (json['connections'] as num?)?.toInt() ?? 0, autoTorrent: json['autoTorrent'] as bool?, @@ -29,6 +39,9 @@ OptsExtraHttp _$OptsExtraHttpFromJson(Map json) => OptsExtraHtt autoExtract: json['autoExtract'] as bool?, archivePassword: json['archivePassword'] as String? ?? '', deleteAfterExtract: json['deleteAfterExtract'] as bool? ?? false, + checksum: json['checksum'] == null + ? null + : ChecksumOption.fromJson(json['checksum'] as Map), ); Map _$OptsExtraHttpToJson(OptsExtraHttp instance) => { @@ -38,4 +51,5 @@ Map _$OptsExtraHttpToJson(OptsExtraHttp instance) => { final _httpHeadersController = AppHttpHeadersController(defaultNames: const ['User-Agent', 'Cookie', 'Referer']); final _trackersController = TextEditingController(); final _archivePasswordController = TextEditingController(); + final _checksumHashController = TextEditingController(); final _formScrollController = ScrollController(); String _httpMethod = 'GET'; @@ -74,6 +75,7 @@ class _CreateTaskWindowPageState extends ConsumerState { Map? _initialLabels; RequestProxyMode _proxyMode = RequestProxyMode.follow; String _proxyScheme = 'http'; + String _checksumAlgorithm = 'sha256'; bool _skipVerifyCert = false; bool? _autoTorrent; bool? _deleteTorrentAfterDownload; @@ -134,6 +136,7 @@ class _CreateTaskWindowPageState extends ConsumerState { _httpHeadersController.dispose(); _trackersController.dispose(); _archivePasswordController.dispose(); + _checksumHashController.dispose(); _formScrollController.dispose(); super.dispose(); } @@ -243,6 +246,12 @@ class _CreateTaskWindowPageState extends ConsumerState { _autoExtract = extra.autoExtract; _archivePasswordController.text = extra.archivePassword; _deleteAfterExtract = extra.deleteAfterExtract; + if (extra.checksum != null) { + if (extra.checksum!.algorithm.isNotEmpty) { + _checksumAlgorithm = extra.checksum!.algorithm; + } + _checksumHashController.text = extra.checksum!.expected; + } } } }); @@ -668,6 +677,32 @@ class _CreateTaskWindowPageState extends ConsumerState { ), ), ], + const SizedBox(height: 14), + AppFormPair( + direction: stacked ? Axis.vertical : Axis.horizontal, + firstLabel: 'Checksum Algorithm', + secondLabel: 'Checksum Hash', + firstFlex: 382, + secondFlex: 618, + first: AppChoiceSegmentedControl( + key: const ValueKey('create-task-checksum-algorithm'), + showIcons: false, + value: _checksumAlgorithm, + buttonKeyPrefix: 'create-task-checksum-algorithm', + alignment: WrapAlignment.start, + options: const [ + AppChoiceOption(value: 'md5', label: 'MD5', icon: Icons.tag), + AppChoiceOption(value: 'sha1', label: 'SHA-1', icon: Icons.tag), + AppChoiceOption(value: 'sha256', label: 'SHA-256', icon: Icons.tag), + ], + onChanged: (value) => setState(() => _checksumAlgorithm = value), + ), + second: _WindowTextField( + key: const ValueKey('create-task-checksum-expected'), + controller: _checksumHashController, + hintText: 'Expected hash (optional)', + ), + ), ] else ...[ AppFormRow( direction: stacked ? Axis.vertical : Axis.horizontal, @@ -1024,6 +1059,13 @@ class _CreateTaskWindowPageState extends ConsumerState { api_options.Options _buildOptions({String? name, List selectFiles = const []}) { final connections = int.tryParse(_connectionsController.text.trim()) ?? 0; + final checksumExpected = _checksumHashController.text.trim(); + final checksum = checksumExpected.isNotEmpty + ? api_options.ChecksumOption( + algorithm: _checksumAlgorithm, + expected: checksumExpected, + ) + : null; return api_options.Options( name: name ?? _renameController.text.trim(), path: _directoryController.text.trim(), @@ -1036,6 +1078,7 @@ class _CreateTaskWindowPageState extends ConsumerState { autoExtract: _autoExtract, archivePassword: _archivePasswordController.text, deleteAfterExtract: _deleteAfterExtract, + checksum: checksum, ).toJson(), ); }