gopeed/pkg/download/extension_extra_test.go

package download

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/GopeedLab/gopeed/pkg/base"
	httpProtocol "github.com/GopeedLab/gopeed/pkg/protocol/http"
)

func TestExtensionExtraMutationHTTP(t *testing.T) {
	for _, resolved := range []bool{false, true} {
		for _, invalid := range []bool{false, true} {
			t.Run(fmt.Sprintf("resolved=%v/invalid=%v", resolved, invalid), func(t *testing.T) {
				var updatedRequests atomic.Int32
				server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					if r.Header.Get("X-New") == "new" {
						updatedRequests.Add(1)
						if r.Header.Get("X-Old") != "" {
							t.Error("extra mutation retained the old header")
						}
					}
					http.ServeContent(w, r, "test.txt", time.Time{}, strings.NewReader("extra mutation payload"))
				}))
				defer server.Close()
				d := NewDownloader(&DownloaderConfig{Storage: NewMemStorage(), StorageDir: t.TempDir()})
				if err := d.Setup(); err != nil {
					t.Fatal(err)
				}
				defer d.Clear()
				extDir := t.TempDir()
				manifest := `{"name":"http-methods","title":"HTTP methods test","version":"0.0.1","scripts":[{"event":"onStart","match":{"labels":["test"]},"entry":"index.js"}]}`
				mutation := `await ctx.task.meta.req.setMethod('GET');
                  await ctx.task.meta.req.setBody('');
                  await ctx.task.meta.req.delHeader('x-old');
                  await ctx.task.meta.req.putHeader('X-New', 'new');`
				if invalid {
					mutation = `try { await ctx.task.meta.req.setHeaders({'X-New':123}); }
                      catch (err) { await ctx.task.meta.req.putLabel('rejected','true'); } ` + mutation
				}
				script := `gopeed.events.onStart(async ctx => { ` + mutation + ` });`
				for name, content := range map[string]string{"manifest.json": manifest, "index.js": script} {
					if err := os.WriteFile(filepath.Join(extDir, name), []byte(content), 0600); err != nil {
						t.Fatal(err)
					}
				}
				if _, err := d.InstallExtensionByFolder(extDir, false); err != nil {
					t.Fatal(err)
				}
				done := make(chan error, 1)
				d.Listener(func(event *Event) {
					if event.Key == EventKeyFinally {
						done <- event.Err
					}
				})
				req := &base.Request{URL: server.URL + "/test.txt", Labels: map[string]string{"test": "true"}, Extra: &httpProtocol.ReqExtra{Header: map[string]string{"X-Old": "old"}}}
				opts := &base.Options{Path: t.TempDir()}
				var id string
				var err error
				if resolved {
					rr, resolveErr := d.Resolve(req)
					if resolveErr != nil {
						t.Fatal(resolveErr)
					}
					id, err = d.CreateWithOptions(rr.ID, opts)
				} else {
					id, err = d.CreateDirect(req, opts)
				}
				if err != nil {
					t.Fatal(err)
				}
				select {
				case err = <-done:
				case <-time.After(10 * time.Second):
					t.Fatal("download timed out")
				}
				if invalid && d.GetTask(id).Meta.Req.Labels["rejected"] != "true" {
					t.Fatal("invalid header was not rejected")
				}
				if err != nil {
					t.Fatal(err)
				}
				got, ok := d.GetTask(id).Meta.Req.Extra.(*httpProtocol.ReqExtra)
				if !ok || got.Header["X-New"] != "new" || len(got.Header) != 1 {
					t.Fatalf("unexpected protocol extra: %#v", d.GetTask(id).Meta.Req.Extra)
				}
				if !resolved && updatedRequests.Load() == 0 {
					t.Fatal("download did not use the extension headers")
				}
				data, err := os.ReadFile(filepath.Join(opts.Path, "test.txt"))
				if err != nil || string(data) != "extra mutation payload" {
					t.Fatalf("download content = %q, error = %v", data, err)
				}
			})
		}
	}
}