diff --git a/pkg/download/downloader.go b/pkg/download/downloader.go index 740b4d255..f49ffeb3a 100644 --- a/pkg/download/downloader.go +++ b/pkg/download/downloader.go @@ -28,6 +28,8 @@ const ( bucketConfig = "config" // downloader extension bucket bucketExtension = "extension" + // downloader extension storage bucket + bucketExtensionStorage = "extension_storage" ) var ( @@ -97,7 +99,7 @@ func NewDownloader(cfg *DownloaderConfig) *Downloader { func (d *Downloader) Setup() error { // setup storage - if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketConfig, bucketExtension}); err != nil { + if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketConfig, bucketExtension, bucketExtensionStorage}); err != nil { return err } // load config from storage diff --git a/pkg/download/extension.go b/pkg/download/extension.go index 4a3eef293..c5622f903 100644 --- a/pkg/download/extension.go +++ b/pkg/download/extension.go @@ -257,6 +257,10 @@ func (d *Downloader) triggerOnResolve(req *base.Request) (res *base.Resource) { } defer scriptFile.Close() ctx.Settings = parseSettings(ext.Settings) + ctx.Storage = &ContextStorage{ + storage: d.storage, + identity: ext.buildIdentity(), + } var scriptBuf []byte scriptBuf, err = io.ReadAll(scriptFile) if err != nil { @@ -507,9 +511,10 @@ func newInstanceLogger(extension *Extension, logger *logger.Logger) *InstanceLog } type Context struct { - Req *base.Request `json:"req"` - Res *base.Resource `json:"res"` - Settings map[string]any `json:"settings"` + Req *base.Request `json:"req"` + Res *base.Resource `json:"res"` + Settings map[string]any `json:"settings"` + Storage *ContextStorage `json:"storage"` } func parseSettings(settings []*Setting) map[string]any { @@ -547,3 +552,50 @@ func tryParse(val any, settingType SettingType) any { return nil } } + +type ContextStorage struct { + storage Storage + identity string +} + +func (s *ContextStorage) Get(key string) any { + raw := s.getRawData() + if v, ok := raw[key]; ok { + return v + } + return nil +} + +func (s *ContextStorage) Set(key string, value string) { + raw := s.getRawData() + raw[key] = value + s.storage.Put(bucketExtensionStorage, s.identity, raw) +} + +func (s *ContextStorage) Remove(key string) { + raw := s.getRawData() + delete(raw, key) + s.storage.Put(bucketExtensionStorage, s.identity, raw) +} + +func (s *ContextStorage) Keys() []string { + raw := s.getRawData() + keys := make([]string, 0) + for k := range raw { + keys = append(keys, k) + } + return keys +} + +func (s *ContextStorage) Clear() { + s.storage.Delete(bucketExtensionStorage, s.identity) +} + +func (s *ContextStorage) getRawData() map[string]string { + var data map[string]string + s.storage.Get(bucketExtensionStorage, s.identity, &data) + if data == nil { + data = make(map[string]string) + } + return data +} diff --git a/pkg/download/extension_test.go b/pkg/download/extension_test.go index a8cb13d3e..f40e24e9d 100644 --- a/pkg/download/extension_test.go +++ b/pkg/download/extension_test.go @@ -170,6 +170,23 @@ func TestDownloader_Extension_Settings(t *testing.T) { }) } +func TestDownloader_ExtensionStorage(t *testing.T) { + setupDownloader(func(downloader *Downloader) { + if _, err := downloader.InstallExtensionByFolder("./testdata/extensions/storage", false); err != nil { + t.Fatal(err) + } + rr, err := downloader.Resolve(&base.Request{ + URL: "https://github.com/test", + }) + if err != nil { + t.Fatal(err) + } + if len(rr.Res.Files) == 1 { + t.Fatal("resolve error") + } + }) +} + func TestDownloader_SwitchExtension(t *testing.T) { setupDownloader(func(downloader *Downloader) { installedExt, err := downloader.InstallExtensionByFolder("./testdata/extensions/basic", false) diff --git a/pkg/download/testdata/extensions/storage/index.js b/pkg/download/testdata/extensions/storage/index.js new file mode 100644 index 000000000..d101e4aa0 --- /dev/null +++ b/pkg/download/testdata/extensions/storage/index.js @@ -0,0 +1,49 @@ +gopeed.events.onResolve(async function (ctx) { + const key = "key" + const value1 = "value1", value2 = JSON.stringify({a: 1, b: "2"}) + + if (ctx.storage.get(key) !== null) { + throw new Error("storage get null error") + } + ctx.storage.remove(key) + if(ctx.storage.keys().length !== 0) { + throw new Error("storage keys null error") + } + + ctx.storage.set(key, value1) + if (ctx.storage.get(key) !== value1) { + throw new Error("storage put1 error") + } + + ctx.storage.set(key, value2) + if (ctx.storage.get(key) !== value2) { + throw new Error("storage put2 error") + } + + if(ctx.storage.keys().length !== 1) { + throw new Error("storage keys error") + } + + ctx.storage.remove(key) + if (ctx.storage.get(key) !== null) { + throw new Error("storage delete error") + } + + ctx.storage.set(key, value1) + ctx.storage.clear() + if (ctx.storage.get(key) !== null) { + throw new Error("storage clear error") + } + + ctx.res = { + name: "test", + files: Array(2).fill(true).map((_, i) => ({ + name: `test-${i}.txt`, + size: 1024, + req: { + url: ctx.req.url + "/" + i, + } + }), + ), + }; +}); diff --git a/pkg/download/testdata/extensions/storage/manifest.json b/pkg/download/testdata/extensions/storage/manifest.json new file mode 100644 index 000000000..ab8ad2bc2 --- /dev/null +++ b/pkg/download/testdata/extensions/storage/manifest.json @@ -0,0 +1,21 @@ +{ + "name": "basic", + "title": "gopeed extension basic test", + "version": "0.0.1", + "scripts": [ + { + "event": "onResolve", + "matches": [ + "*://github.com/*" + ], + "entry": "index.js" + } + ], + "settings": [ + { + "name": "ua", + "title": "User-Agent", + "type": "string" + } + ] +} \ No newline at end of file