feat: add gblob scheme support for extension (#1358)

pkg/download/engine/engine.go

@@ -3,10 +3,12 @@ package engine
 import (
 	_ "embed"
 	"errors"
+
 	"github.com/GopeedLab/gopeed/pkg/base"
 	gojaerror "github.com/GopeedLab/gopeed/pkg/download/engine/inject/error"
 	"github.com/GopeedLab/gopeed/pkg/download/engine/inject/file"
 	"github.com/GopeedLab/gopeed/pkg/download/engine/inject/formdata"
+	"github.com/GopeedLab/gopeed/pkg/download/engine/inject/stream"
 	"github.com/GopeedLab/gopeed/pkg/download/engine/inject/vm"
 	"github.com/GopeedLab/gopeed/pkg/download/engine/inject/xhr"
 	"github.com/dop251/goja"
@@ -27,90 +29,62 @@ type Engine struct {
 // RunString executes the script and returns the go type value
 // if script result is promise, it will be resolved
 func (e *Engine) RunString(script string) (value any, err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			err = r.(error)
-		}
-	}()
-
-	var result goja.Value
-	e.loop.Run(func(runtime *goja.Runtime) {
-		result, err = runtime.RunString(script)
-		if err == nil {
-			go e.await(result)
-		}
+	return e.runOnLoop(func(runtime *goja.Runtime) (goja.Value, error) {
+		return runtime.RunString(script)
 	})
-	if err != nil {
-		return
-	}
-	return resolveResult(result)
 }
 
 // CallFunction calls the function and returns the go type value
 // if function result is promise, it will be resolved
 func (e *Engine) CallFunction(fn goja.Callable, args ...any) (value any, err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			err = r.(error)
-		}
-	}()
-
-	var result goja.Value
-	e.loop.Run(func(runtime *goja.Runtime) {
+	return e.runOnLoop(func(runtime *goja.Runtime) (goja.Value, error) {
 		if args == nil {
-			result, err = fn(nil)
-		} else {
-			var jsArgs []goja.Value
-			for _, arg := range args {
-				jsArgs = append(jsArgs, runtime.ToValue(arg))
-			}
-			result, err = fn(nil, jsArgs...)
+			return fn(nil)
 		}
-		if err == nil {
-			go e.await(result)
+		var jsArgs []goja.Value
+		for _, arg := range args {
+			jsArgs = append(jsArgs, runtime.ToValue(arg))
 		}
+		return fn(nil, jsArgs...)
 	})
-	if err != nil {
-		return
-	}
-	return resolveResult(result)
 }
 
-// loop.Run will hang if the script result has a non-stop code, such as setInterval.
-// This method will stop the event loop when the promise result is resolved.
-func (e *Engine) await(value any) {
-	if value == nil {
-		return
+func (e *Engine) runOnLoop(fn func(runtime *goja.Runtime) (goja.Value, error)) (any, error) {
+	type result struct {
+		value goja.Value
+		err   error
 	}
-
-	if v, ok := value.(goja.Value); ok {
-		// if result is promise, wait for it to be resolved
-		if p, ok := v.Export().(*goja.Promise); ok {
-			if p.State() != goja.PromiseStatePending {
-				return
-			}
-
-			// check promise state every 100 milliseconds, until it is resolved
-			for {
-				time.Sleep(time.Millisecond * 100)
-				if p.State() == goja.PromiseStatePending {
-					continue
+	ch := make(chan result, 1)
+	ok := e.loop.RunOnLoop(func(runtime *goja.Runtime) {
+		defer func() {
+			if r := recover(); r != nil {
+				if err, ok := r.(error); ok {
+					ch <- result{err: err}
+					return
 				}
-				break
+				ch <- result{err: errors.New("panic")}
 			}
-
-			// stop the event loop
-			e.loop.StopNoWait()
-		}
+		}()
+		value, err := fn(runtime)
+		ch <- result{value: value, err: err}
+	})
+	if !ok {
+		return nil, errors.New("engine loop terminated")
 	}
+	res := <-ch
+	if res.err != nil {
+		return nil, res.err
+	}
+	return resolveResult(res.value)
 }
 
 func (e *Engine) Close() {
-	e.loop.StopNoWait()
+	e.loop.Terminate()
 }
 
 type Config struct {
-	ProxyConfig *base.DownloaderProxyConfig
+	ProxyConfig  *base.DownloaderProxyConfig
+	StreamConfig *stream.Config
 }
 
 func NewEngine(cfg *Config) *Engine {
@@ -121,7 +95,10 @@ func NewEngine(cfg *Config) *Engine {
 	engine := &Engine{
 		loop: loop,
 	}
-	loop.Run(func(runtime *goja.Runtime) {
+	loop.Start()
+	done := make(chan struct{})
+	loop.RunOnLoop(func(runtime *goja.Runtime) {
+		defer close(done)
 		engine.Runtime = runtime
 		runtime.SetFieldNameMapper(goja.TagFieldNameMapper("json", true))
 		vm.Enable(runtime)
@@ -153,8 +130,12 @@ func NewEngine(cfg *Config) *Engine {
 		if _, err := runtime.RunString("global.location = new URL('http://localhost');"); err != nil {
 			return
 		}
+		if err := stream.Enable(runtime, cfg.StreamConfig); err != nil {
+			return
+		}
 		return
 	})
+	<-done
 	return engine
 }
 
@@ -169,6 +150,9 @@ func resolveResult(value goja.Value) (any, error) {
 	switch export.(type) {
 	case *goja.Promise:
 		p := export.(*goja.Promise)
+		for p.State() == goja.PromiseStatePending {
+			time.Sleep(time.Millisecond * 10)
+		}
 		switch p.State() {
 		case goja.PromiseStatePending:
 			return nil, nil

pkg/download/engine/inject/stream/module.go

@@ -0,0 +1,92 @@
+package stream
+
+import (
+	_ "embed"
+	"fmt"
+
+	"github.com/dop251/goja"
+)
+
+//go:embed stream.js
+var script string
+
+type Config struct {
+	CreateBlobObjectURL           func(data []byte, contentType string) (string, error)
+	CreateWritableStreamObjectURL func() (string, error)
+	WriteWritableStreamObjectURL  func(url string, data any) error
+	CloseWritableStreamObjectURL  func(url string) error
+	AbortWritableStreamObjectURL  func(url string, reason string) error
+	RevokeObjectURL               func(url string) error
+}
+
+func Enable(runtime *goja.Runtime, cfg *Config) error {
+	if cfg == nil {
+		cfg = &Config{}
+	}
+	if err := runtime.Set("__gopeed_create_blob_object_url", func(data []byte, contentType string) string {
+		if cfg.CreateBlobObjectURL == nil {
+			panic(fmt.Errorf("gblob blob object url handler not configured"))
+		}
+		url, err := cfg.CreateBlobObjectURL(data, contentType)
+		if err != nil {
+			panic(err)
+		}
+		return url
+	}); err != nil {
+		return err
+	}
+	if err := runtime.Set("__gopeed_create_writable_stream_object_url", func() string {
+		if cfg.CreateWritableStreamObjectURL == nil {
+			panic(fmt.Errorf("gblob writable stream object url handler not configured"))
+		}
+		url, err := cfg.CreateWritableStreamObjectURL()
+		if err != nil {
+			panic(err)
+		}
+		return url
+	}); err != nil {
+		return err
+	}
+	if err := runtime.Set("__gopeed_write_writable_stream_object_url", func(url string, data any) {
+		if cfg.WriteWritableStreamObjectURL == nil {
+			panic(fmt.Errorf("gblob writable stream write handler not configured"))
+		}
+		if err := cfg.WriteWritableStreamObjectURL(url, data); err != nil {
+			panic(err)
+		}
+	}); err != nil {
+		return err
+	}
+	if err := runtime.Set("__gopeed_close_writable_stream_object_url", func(url string) {
+		if cfg.CloseWritableStreamObjectURL == nil {
+			panic(fmt.Errorf("gblob writable stream close handler not configured"))
+		}
+		if err := cfg.CloseWritableStreamObjectURL(url); err != nil {
+			panic(err)
+		}
+	}); err != nil {
+		return err
+	}
+	if err := runtime.Set("__gopeed_abort_writable_stream_object_url", func(url string, reason string) {
+		if cfg.AbortWritableStreamObjectURL == nil {
+			panic(fmt.Errorf("gblob writable stream abort handler not configured"))
+		}
+		if err := cfg.AbortWritableStreamObjectURL(url, reason); err != nil {
+			panic(err)
+		}
+	}); err != nil {
+		return err
+	}
+	if err := runtime.Set("__gopeed_revoke_object_url", func(url string) {
+		if cfg.RevokeObjectURL == nil {
+			panic(fmt.Errorf("gblob revoke handler not configured"))
+		}
+		if err := cfg.RevokeObjectURL(url); err != nil {
+			panic(err)
+		}
+	}); err != nil {
+		return err
+	}
+	_, err := runtime.RunString(script)
+	return err
+}

pkg/download/engine/inject/stream/stream.js

@@ -0,0 +1,341 @@
+(function () {
+  const createBlobObjectURL = globalThis.__gopeed_create_blob_object_url;
+  const createWritableObjectURL = globalThis.__gopeed_create_writable_stream_object_url;
+  const writeWritableObjectURL = globalThis.__gopeed_write_writable_stream_object_url;
+  const closeWritableObjectURL = globalThis.__gopeed_close_writable_stream_object_url;
+  const abortWritableObjectURL = globalThis.__gopeed_abort_writable_stream_object_url;
+  const revokeObjectURL = globalThis.__gopeed_revoke_object_url;
+
+  if (typeof globalThis.ReadableStream === "undefined") {
+    class ReadableStreamDefaultController {
+      constructor(stream) {
+        this._stream = stream;
+      }
+
+      enqueue(chunk) {
+        this._stream._enqueue(chunk);
+      }
+
+      close() {
+        this._stream._close();
+      }
+
+      error(err) {
+        this._stream._error(err);
+      }
+    }
+
+    class ReadableStreamDefaultReader {
+      constructor(stream) {
+        this._stream = stream;
+      }
+
+      read() {
+        return this._stream._read();
+      }
+
+      releaseLock() {
+        if (this._stream) {
+          this._stream._reader = null;
+          this._stream.locked = false;
+          this._stream = null;
+        }
+      }
+    }
+
+    class ReadableStream {
+      constructor(source = {}) {
+        this._queue = [];
+        this._waiters = [];
+        this._closed = false;
+        this._errored = null;
+        this._reader = null;
+        this.locked = false;
+        const controller = new ReadableStreamDefaultController(this);
+        if (typeof source.start === "function") {
+          source.start(controller);
+        }
+      }
+
+      _enqueue(chunk) {
+        if (this._closed || this._errored) {
+          return;
+        }
+        if (this._waiters.length > 0) {
+          const waiter = this._waiters.shift();
+          waiter.resolve({ done: false, value: chunk });
+          return;
+        }
+        this._queue.push(chunk);
+      }
+
+      _close() {
+        this._closed = true;
+        while (this._waiters.length > 0) {
+          this._waiters.shift().resolve({ done: true, value: undefined });
+        }
+      }
+
+      _error(err) {
+        this._errored = err || new Error("ReadableStream error");
+        while (this._waiters.length > 0) {
+          this._waiters.shift().reject(this._errored);
+        }
+      }
+
+      _read() {
+        if (this._errored) {
+          return Promise.reject(this._errored);
+        }
+        if (this._queue.length > 0) {
+          return Promise.resolve({ done: false, value: this._queue.shift() });
+        }
+        if (this._closed) {
+          return Promise.resolve({ done: true, value: undefined });
+        }
+        return new Promise((resolve, reject) => {
+          this._waiters.push({ resolve, reject });
+        });
+      }
+
+      getReader() {
+        if (this.locked) {
+          throw new TypeError("ReadableStream is locked");
+        }
+        this.locked = true;
+        this._reader = new ReadableStreamDefaultReader(this);
+        return this._reader;
+      }
+
+      async pipeTo(dest) {
+        const reader = this.getReader();
+        const writer = dest.getWriter();
+        try {
+          while (true) {
+            const { done, value } = await reader.read();
+            if (done) {
+              await writer.close();
+              break;
+            }
+            await writer.write(value);
+          }
+        } catch (err) {
+          await writer.abort(err);
+          throw err;
+        } finally {
+          reader.releaseLock();
+          writer.releaseLock();
+        }
+      }
+    }
+
+    globalThis.ReadableStream = ReadableStream;
+  }
+
+  if (typeof globalThis.WritableStream === "undefined") {
+    class WritableStreamDefaultWriter {
+      constructor(stream) {
+        if (stream.locked) {
+          throw new TypeError("WritableStream is locked");
+        }
+        this._stream = stream;
+        this._stream.locked = true;
+      }
+
+      write(chunk) {
+        return this._stream._write(chunk);
+      }
+
+      close() {
+        return this._stream._close();
+      }
+
+      abort(reason) {
+        return this._stream._abort(reason);
+      }
+
+      releaseLock() {
+        if (this._stream) {
+          this._stream.locked = false;
+          this._stream = null;
+        }
+      }
+    }
+
+    class WritableStream {
+      constructor(sink = {}) {
+        this._sink = sink;
+        this._observers = [];
+        this._state = "writable";
+        this.locked = false;
+      }
+
+      _targets() {
+        return [this._sink, ...this._observers];
+      }
+
+      _write(chunk) {
+        if (this._state !== "writable") {
+          return Promise.reject(new TypeError("WritableStream is not writable"));
+        }
+        let chain = Promise.resolve();
+        for (const target of this._targets()) {
+          if (target && typeof target.write === "function") {
+            chain = chain.then(() => target.write(chunk));
+          }
+        }
+        return chain;
+      }
+
+      _close() {
+        if (this._state !== "writable") {
+          return Promise.reject(new TypeError("WritableStream is not writable"));
+        }
+        this._state = "closed";
+        let chain = Promise.resolve();
+        for (const target of this._targets()) {
+          if (target && typeof target.close === "function") {
+            chain = chain.then(() => target.close());
+          }
+        }
+        return chain;
+      }
+
+      _abort(reason) {
+        if (this._state === "errored") {
+          return Promise.resolve();
+        }
+        this._state = "errored";
+        let chain = Promise.resolve();
+        for (const target of this._targets()) {
+          if (target && typeof target.abort === "function") {
+            chain = chain.then(() => target.abort(reason));
+          }
+        }
+        return chain;
+      }
+
+      _addObserver(observer) {
+        this._observers.push(observer);
+      }
+
+      getWriter() {
+        return new WritableStreamDefaultWriter(this);
+      }
+    }
+
+    globalThis.WritableStream = WritableStream;
+  }
+
+  if (typeof globalThis.TransformStream === "undefined") {
+    class TransformStream {
+      constructor(transformer = {}) {
+        let readableController;
+        this.readable = new ReadableStream({
+          start(controller) {
+            readableController = controller;
+          }
+        });
+        this.writable = new WritableStream({
+          async write(chunk) {
+            if (typeof transformer.transform === "function") {
+              await transformer.transform(chunk, readableController);
+              return;
+            }
+            readableController.enqueue(chunk);
+          },
+          async close() {
+            if (typeof transformer.flush === "function") {
+              await transformer.flush(readableController);
+            }
+            readableController.close();
+          },
+          async abort(reason) {
+            if (typeof transformer.abort === "function") {
+              await transformer.abort(reason);
+            }
+            readableController.error(reason || new Error("TransformStream aborted"));
+          }
+        });
+      }
+    }
+
+    globalThis.TransformStream = TransformStream;
+  }
+
+  if (typeof globalThis.ByteLengthQueuingStrategy === "undefined") {
+    globalThis.ByteLengthQueuingStrategy = class ByteLengthQueuingStrategy {
+      constructor({ highWaterMark }) {
+        this.highWaterMark = highWaterMark;
+      }
+
+      size(chunk) {
+        if (typeof chunk === "string") {
+          return chunk.length;
+        }
+        if (chunk && typeof chunk.byteLength === "number") {
+          return chunk.byteLength;
+        }
+        return 1;
+      }
+    };
+  }
+
+  if (typeof globalThis.CountQueuingStrategy === "undefined") {
+    globalThis.CountQueuingStrategy = class CountQueuingStrategy {
+      constructor({ highWaterMark }) {
+        this.highWaterMark = highWaterMark;
+      }
+
+      size() {
+        return 1;
+      }
+    };
+  }
+
+  const originalCreateObjectURL = typeof URL.createObjectURL === "function"
+    ? URL.createObjectURL.bind(URL)
+    : null;
+  const originalRevokeObjectURL = typeof URL.revokeObjectURL === "function"
+    ? URL.revokeObjectURL.bind(URL)
+    : null;
+
+  URL.createObjectURL = function (value) {
+    if (value instanceof Blob) {
+      return createBlobObjectURL(value._buffer, value.type || "");
+    }
+    if (value instanceof WritableStream) {
+      if (value.__gopeedObjectURL) {
+        return value.__gopeedObjectURL;
+      }
+      const url = createWritableObjectURL();
+      value._addObserver({
+        write(chunk) {
+          return Promise.resolve(writeWritableObjectURL(url, chunk));
+        },
+        close() {
+          return Promise.resolve(closeWritableObjectURL(url));
+        },
+        abort(reason) {
+          return Promise.resolve(abortWritableObjectURL(url, reason == null ? "" : String(reason)));
+        }
+      });
+      value.__gopeedObjectURL = url;
+      return url;
+    }
+    if (originalCreateObjectURL) {
+      return originalCreateObjectURL(value);
+    }
+    throw new TypeError("Unsupported object type");
+  };
+
+  URL.revokeObjectURL = function (url) {
+    if (typeof url === "string" && url.indexOf("gblob:") === 0) {
+      revokeObjectURL(url);
+      return;
+    }
+    if (originalRevokeObjectURL) {
+      originalRevokeObjectURL(url);
+    }
+  };
+})();