feat: support storage state on extension (#235)

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

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

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)

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

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"
+    }
+  ]
+}