feat: enhance extension API (#1413)

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")

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)
 	}
 }
 

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},
 	})

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)
+		}
 	})
 }
 

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) => ({

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();
 });

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();
 });

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;
   }
 });

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();
 });
-

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');
 });
-