diff --git a/pkg/base/model.go b/pkg/base/model.go index c4a9d7d4..b99d5d53 100644 --- a/pkg/base/model.go +++ b/pkg/base/model.go @@ -25,6 +25,24 @@ type Request struct { SkipVerifyCert bool `json:"skipVerifyCert"` } +// SetLabels replaces all labels on the request. +func (r *Request) SetLabels(labels map[string]string) { + r.Labels = labels +} + +// PutLabel sets a label on the request. +func (r *Request) PutLabel(key, value string) { + if r.Labels == nil { + r.Labels = make(map[string]string) + } + r.Labels[key] = value +} + +// DelLabel deletes a label from the request. +func (r *Request) DelLabel(key string) { + delete(r.Labels, key) +} + func (r *Request) Validate() error { if r.URL == "" { return fmt.Errorf("invalid request url") diff --git a/pkg/download/downloader_test.go b/pkg/download/downloader_test.go index f7eb9974..0f98890d 100644 --- a/pkg/download/downloader_test.go +++ b/pkg/download/downloader_test.go @@ -42,35 +42,42 @@ type generationTestManager struct { pauseOnce sync.Once } -func TestExtensionTaskControlMethodsOnlyExistOnError(t *testing.T) { +func TestExtensionTaskMethodsMatchEventCapabilities(t *testing.T) { taskType := reflect.TypeOf((*Task)(nil)) if _, ok := taskType.MethodByName("Continue"); ok { t.Fatal("regular task unexpectedly exposes Continue") } - if _, ok := taskType.MethodByName("Pause"); ok { - t.Fatal("regular task unexpectedly exposes Pause") + if _, ok := taskType.MethodByName("SetUrl"); ok { + t.Fatal("regular task unexpectedly exposes SetUrl") } - errorTaskType := reflect.TypeOf((*ExtensionTask)(nil)) + extensionTaskType := reflect.TypeOf((*ExtensionTask)(nil)) + if _, ok := extensionTaskType.MethodByName("SetUrl"); !ok { + t.Fatal("extension task does not expose SetUrl") + } + if _, ok := extensionTaskType.MethodByName("Continue"); ok { + t.Fatal("non-error extension task unexpectedly exposes Continue") + } + + errorTaskType := reflect.TypeOf((*OnErrorExtensionTask)(nil)) + if _, ok := errorTaskType.MethodByName("SetUrl"); !ok { + t.Fatal("onError task does not expose SetUrl") + } if _, ok := errorTaskType.MethodByName("Continue"); !ok { t.Fatal("onError task does not expose Continue") } - if _, ok := errorTaskType.MethodByName("Pause"); ok { - t.Fatal("onError task unexpectedly exposes Pause") - } - for name, contextType := range map[string]reflect.Type{ - "onStart": reflect.TypeOf(OnStartContext{}), - "onDone": reflect.TypeOf(OnDoneContext{}), - } { - field, _ := contextType.FieldByName("Task") - if field.Type != taskType { - t.Fatalf("%s unexpectedly injects a control wrapper: %v", name, field.Type) - } + startTaskField, _ := reflect.TypeOf(OnStartContext{}).FieldByName("Task") + if startTaskField.Type != extensionTaskType { + t.Fatalf("onStart does not inject the request mutation wrapper: %v", startTaskField.Type) + } + doneTaskField, _ := reflect.TypeOf(OnDoneContext{}).FieldByName("Task") + if doneTaskField.Type != taskType { + t.Fatalf("onDone unexpectedly injects a mutation wrapper: %v", doneTaskField.Type) } errorTaskField, _ := reflect.TypeOf(OnErrorContext{}).FieldByName("Task") if errorTaskField.Type != errorTaskType { - t.Fatalf("onError does not inject the Continue wrapper: %v", errorTaskField.Type) + t.Fatalf("onError does not inject the recovery wrapper: %v", errorTaskField.Type) } } diff --git a/pkg/download/extension.go b/pkg/download/extension.go index a38038c4..c9bb7832 100644 --- a/pkg/download/extension.go +++ b/pkg/download/extension.go @@ -705,12 +705,12 @@ type OnResolveContext struct { } type OnStartContext struct { - Task *Task `json:"task"` + Task *ExtensionTask `json:"task"` } type OnErrorContext struct { - Task *ExtensionTask `json:"task"` - Error error `json:"error"` + Task *OnErrorExtensionTask `json:"task"` + Error error `json:"error"` } type OnDoneContext struct { @@ -721,9 +721,15 @@ type OnDoneContext struct { // Avoid extension scripts modifying task directly, use ExtensionTask to encapsulate task, // only some fields can be modified, such as request info. type ExtensionTask struct { + *Task +} + +// OnErrorExtensionTask adds error-recovery controls to ExtensionTask. +// Continue is intentionally only exposed to onError handlers. +type OnErrorExtensionTask struct { download *Downloader - *Task + *ExtensionTask } func cloneExtensionTask(task *Task) *Task { @@ -732,14 +738,18 @@ func cloneExtensionTask(task *Task) *Task { return newTask } -func newOnStartExtensionTask(task *Task) *Task { - return cloneExtensionTask(task) +func newExtensionTask(task *Task) *ExtensionTask { + return &ExtensionTask{Task: cloneExtensionTask(task)} } -func newOnErrorExtensionTask(download *Downloader, task *Task) *ExtensionTask { - return &ExtensionTask{ - download: download, - Task: cloneExtensionTask(task), +func newOnStartExtensionTask(task *Task) *ExtensionTask { + return newExtensionTask(task) +} + +func newOnErrorExtensionTask(download *Downloader, task *Task) *OnErrorExtensionTask { + return &OnErrorExtensionTask{ + download: download, + ExtensionTask: newExtensionTask(task), } } @@ -747,7 +757,12 @@ func newOnDoneExtensionTask(task *Task) *Task { return cloneExtensionTask(task) } -func (t *ExtensionTask) Continue() error { +// SetUrl replaces the task request URL. +func (t *ExtensionTask) SetUrl(url string) { + t.Meta.Req.URL = url +} + +func (t *OnErrorExtensionTask) Continue() error { return t.download.Continue(&TaskFilter{ IDs: []string{t.ID}, }) diff --git a/pkg/download/extension_test.go b/pkg/download/extension_test.go index 0cdb9bd6..fd63d9a5 100644 --- a/pkg/download/extension_test.go +++ b/pkg/download/extension_test.go @@ -26,15 +26,26 @@ func TestDownloader_InstallExtensionByFolder(t *testing.T) { if _, err := downloader.InstallExtensionByFolder("./testdata/extensions/basic", false); err != nil { t.Fatal(err) } - rr, err := downloader.Resolve(&base.Request{ - URL: "https://github.com/test", - }, nil) + req := &base.Request{ + URL: "https://github.com/test", + Labels: map[string]string{"original": "true"}, + } + rr, err := downloader.Resolve(req, nil) if err != nil { t.Fatal(err) } if len(rr.Res.Files) == 1 { t.Fatal("resolve error") } + if req.Labels["replaced"] != "true" || req.Labels["modified"] != "true" { + t.Fatalf("request label methods did not update labels: %#v", req.Labels) + } + if _, ok := req.Labels["original"]; ok { + t.Fatalf("setLabels did not replace existing labels: %#v", req.Labels) + } + if _, ok := req.Labels["removed"]; ok { + t.Fatalf("delLabel did not remove label: %#v", req.Labels) + } }) } diff --git a/pkg/download/testdata/extensions/basic/index.js b/pkg/download/testdata/extensions/basic/index.js index 45fb095a..c4751dbb 100644 --- a/pkg/download/testdata/extensions/basic/index.js +++ b/pkg/download/testdata/extensions/basic/index.js @@ -1,4 +1,7 @@ gopeed.events.onResolve(async function (ctx) { + ctx.req.setLabels({ replaced: "true", removed: "true" }); + ctx.req.putLabel("modified", "true"); + ctx.req.delLabel("removed"); ctx.res = { name: "test", files: Array(2).fill(true).map((_, i) => ({ diff --git a/pkg/download/testdata/extensions/blob_pause_rebuild/index.js b/pkg/download/testdata/extensions/blob_pause_rebuild/index.js index b42a3d91..93263b0c 100644 --- a/pkg/download/testdata/extensions/blob_pause_rebuild/index.js +++ b/pkg/download/testdata/extensions/blob_pause_rebuild/index.js @@ -132,9 +132,8 @@ gopeed.events.onError(async function (ctx) { return; } - req.url = createPayloadURL(source, false); - req.labels = req.labels || {}; - req.labels.started = 'true'; - req.labels.rebuilt = 'true'; + ctx.task.setUrl(createPayloadURL(source, false)); + req.putLabel('started', 'true'); + req.putLabel('rebuilt', 'true'); ctx.task.continue(); }); diff --git a/pkg/download/testdata/extensions/blob_recover/index.js b/pkg/download/testdata/extensions/blob_recover/index.js index 9d331fd0..911ae199 100644 --- a/pkg/download/testdata/extensions/blob_recover/index.js +++ b/pkg/download/testdata/extensions/blob_recover/index.js @@ -40,7 +40,7 @@ gopeed.events.onError(async function (ctx) { return; } - req.labels.recovered = "true"; - req.url = gopeed.runtime.blob.createObjectURL(new Blob(["ok\n"], { type: "text/plain" })); + req.putLabel("recovered", "true"); + ctx.task.setUrl(gopeed.runtime.blob.createObjectURL(new Blob(["ok\n"], { type: "text/plain" }))); ctx.task.continue(); }); diff --git a/pkg/download/testdata/extensions/blob_restart/index.js b/pkg/download/testdata/extensions/blob_restart/index.js index 834e629a..4a7eb240 100644 --- a/pkg/download/testdata/extensions/blob_restart/index.js +++ b/pkg/download/testdata/extensions/blob_restart/index.js @@ -41,18 +41,17 @@ gopeed.events.onError(async function (ctx) { if (!req || !req.rawUrl || !req.url || !req.url.includes('/__blob/')) { return; } - req.labels = req.labels || {}; if (req.labels.rebuilt === 'true') { return; } try { - req.url = createRangeBlobUrl(req.rawUrl); - req.labels.started = 'true'; - req.labels.rebuilt = 'true'; + ctx.task.setUrl(createRangeBlobUrl(req.rawUrl)); + req.putLabel('started', 'true'); + req.putLabel('rebuilt', 'true'); ctx.task.continue(); } catch (error) { - req.labels.rebuildError = String(error); + req.putLabel('rebuildError', String(error)); throw error; } }); diff --git a/pkg/download/testdata/extensions/on_error/index.js b/pkg/download/testdata/extensions/on_error/index.js index 59e0abac..c8fa26b2 100644 --- a/pkg/download/testdata/extensions/on_error/index.js +++ b/pkg/download/testdata/extensions/on_error/index.js @@ -1,7 +1,6 @@ gopeed.events.onError(async function (ctx) { gopeed.logger.info("url", ctx.task.meta.req.url); gopeed.logger.info("error", ctx.error); - ctx.task.meta.req.url = "https://github.com"; + ctx.task.setUrl("https://github.com"); ctx.task.continue(); }); - diff --git a/pkg/download/testdata/extensions/on_start/index.js b/pkg/download/testdata/extensions/on_start/index.js index e41109a5..671c99e9 100644 --- a/pkg/download/testdata/extensions/on_start/index.js +++ b/pkg/download/testdata/extensions/on_start/index.js @@ -1,6 +1,5 @@ gopeed.events.onStart(async function (ctx) { gopeed.logger.info("url", ctx.task.meta.req.url); - ctx.task.meta.req.url = "https://github.com"; - ctx.task.meta.req.labels['modified'] = 'true'; + ctx.task.setUrl("https://github.com"); + ctx.task.meta.req.putLabel('modified', 'true'); }); -