Skip to content
18 changes: 18 additions & 0 deletions pkg/base/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
39 changes: 23 additions & 16 deletions pkg/download/downloader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down
37 changes: 26 additions & 11 deletions pkg/download/extension.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -732,22 +738,31 @@ 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 newOnStartExtensionTask(task *Task) *ExtensionTask {
return newExtensionTask(task)
}

func newOnErrorExtensionTask(download *Downloader, task *Task) *ExtensionTask {
return &ExtensionTask{
download: download,
Task: cloneExtensionTask(task),
func newOnErrorExtensionTask(download *Downloader, task *Task) *OnErrorExtensionTask {
return &OnErrorExtensionTask{
download: download,
ExtensionTask: newExtensionTask(task),
}
}

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},
})
Expand Down
17 changes: 14 additions & 3 deletions pkg/download/extension_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
})
}

Expand Down
3 changes: 3 additions & 0 deletions pkg/download/testdata/extensions/basic/index.js
Original file line number Diff line number Diff line change
@@ -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) => ({
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
});
4 changes: 2 additions & 2 deletions pkg/download/testdata/extensions/blob_recover/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -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();
});
9 changes: 4 additions & 5 deletions pkg/download/testdata/extensions/blob_restart/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
});
3 changes: 1 addition & 2 deletions pkg/download/testdata/extensions/on_error/index.js
Original file line number Diff line number Diff line change
@@ -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();
});

5 changes: 2 additions & 3 deletions pkg/download/testdata/extensions/on_start/index.js
Original file line number Diff line number Diff line change
@@ -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');
});

Loading