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

package stream

import (
	"context"
	_ "embed"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"

	"github.com/dop251/goja"
	"github.com/dop251/goja_nodejs/eventloop"
)

//go:embed stream.js
var script string

type Config struct {
	CreateObjectURL func(opts *ObjectURLOptions, open ObjectURLOpener) (string, error)
	RevokeObjectURL func(url string) error
}

type ObjectURLOptions struct {
	ContentType string
	Size        int64
	Range       bool
}

type ObjectURLOpenRequest struct {
	Offset int64
	End    int64
}

type ObjectURLOpener func(ctx context.Context, req ObjectURLOpenRequest) (io.ReadCloser, error)

func Enable(runtime *goja.Runtime, loop *eventloop.EventLoop, cfg *Config) error {
	if cfg == nil {
		cfg = &Config{}
	}
	blobPipes := newBlobPipeRegistry()
	if err := runtime.Set("__gopeed_create_blob_object_url", func(call goja.FunctionCall) goja.Value {
		if cfg.CreateObjectURL == nil {
			panic(runtime.NewGoError(fmt.Errorf("blob object url handler not configured")))
		}
		openValue := call.Argument(0)
		if _, ok := goja.AssertFunction(openValue); !ok {
			panic(runtime.NewGoError(fmt.Errorf("blob opener must be callable")))
		}
		optsObj := call.Argument(1).ToObject(runtime)
		size := int64(0)
		contentType := ""
		rangeEnabled := false
		if optsObj != nil {
			size = optsObj.Get("size").ToInteger()
			contentType = optsObj.Get("contentType").String()
			rangeEnabled = optsObj.Get("range").ToBoolean()
		}
		if size < 0 {
			size = 0
		}
		opts := &ObjectURLOptions{
			ContentType: contentType,
			Size:        size,
			Range:       rangeEnabled,
		}
		url, err := cfg.CreateObjectURL(opts, func(ctx context.Context, req ObjectURLOpenRequest) (io.ReadCloser, error) {
			return openBlobObjectURLReader(ctx, loop, blobPipes, openValue, req)
		})
		if err != nil {
			panic(runtime.NewGoError(err))
		}
		return runtime.ToValue(url)
	}); err != nil {
		return err
	}
	if err := runtime.Set("__gopeed_revoke_object_url", func(url string) {
		if cfg.RevokeObjectURL == nil {
			panic(runtime.NewGoError(fmt.Errorf("blob revoke handler not configured")))
		}
		if err := cfg.RevokeObjectURL(url); err != nil {
			panic(runtime.NewGoError(err))
		}
	}); err != nil {
		return err
	}
	if err := runtime.Set("__gopeed_blob_pipe_chunk", func(call goja.FunctionCall) goja.Value {
		id := call.Argument(0).String()
		ok := blobPipes.Push(id, exportBytes(call.Argument(1)))
		return runtime.ToValue(ok)
	}); err != nil {
		return err
	}
	if err := runtime.Set("__gopeed_blob_pipe_close", func(id string) {
		blobPipes.Close(id, nil)
	}); err != nil {
		return err
	}
	if err := runtime.Set("__gopeed_blob_pipe_error", func(id string, message string) {
		blobPipes.Close(id, errors.New(message))
	}); err != nil {
		return err
	}
	_, err := runtime.RunString(script)
	return err
}

func openBlobObjectURLReader(ctx context.Context, loop *eventloop.EventLoop, pipes *blobPipeRegistry, openValue goja.Value, req ObjectURLOpenRequest) (io.ReadCloser, error) {
	if pipes == nil {
		return nil, fmt.Errorf("blob pipe registry not configured")
	}
	pipe := newBlobPipe(ctx, loop, pipes)
	pipes.Add(pipe)
	_, err := runOnLoop(loop, func(runtime *goja.Runtime) (goja.Value, error) {
		fnVal := runtime.Get("__gopeed_blob_pipe_source")
		fn, ok := goja.AssertFunction(fnVal)
		if !ok {
			return nil, fmt.Errorf("blob pipe helper is not callable")
		}
		request := map[string]any{
			"offset": req.Offset,
			"end":    req.End,
		}
		return fn(nil, openValue, runtime.ToValue(request), runtime.ToValue(pipe.id))
	})
	if err != nil {
		pipe.Close()
		return nil, err
	}
	return pipe, nil
}

type blobPipeRegistry struct {
	mu    sync.Mutex
	pipes map[string]*blobPipe
	next  atomic.Uint64
}

func newBlobPipeRegistry() *blobPipeRegistry {
	return &blobPipeRegistry{
		pipes: map[string]*blobPipe{},
	}
}

func (r *blobPipeRegistry) Add(pipe *blobPipe) {
	r.mu.Lock()
	r.pipes[pipe.id] = pipe
	r.mu.Unlock()
}

func (r *blobPipeRegistry) Remove(id string) {
	r.mu.Lock()
	delete(r.pipes, id)
	r.mu.Unlock()
}

func (r *blobPipeRegistry) Push(id string, chunk []byte) bool {
	r.mu.Lock()
	pipe := r.pipes[id]
	r.mu.Unlock()
	if pipe == nil {
		return false
	}
	return pipe.Push(chunk)
}

func (r *blobPipeRegistry) Close(id string, err error) {
	r.mu.Lock()
	pipe := r.pipes[id]
	r.mu.Unlock()
	if pipe != nil {
		pipe.CloseWithError(err)
	}
}

func (r *blobPipeRegistry) NewID() string {
	return fmt.Sprintf("pipe-%d", r.next.Add(1))
}

type blobPipeItem struct {
	chunk []byte
	err   error
	done  bool
}

type blobPipe struct {
	id       string
	loop     *eventloop.EventLoop
	registry *blobPipeRegistry
	ctx      context.Context
	cancel   context.CancelFunc
	ch       chan blobPipeItem

	mu     sync.Mutex
	buf    []byte
	closed bool
	once   sync.Once
}

func newBlobPipe(ctx context.Context, loop *eventloop.EventLoop, registry *blobPipeRegistry) *blobPipe {
	pipeCtx, cancel := context.WithCancel(ctx)
	return &blobPipe{
		id:       registry.NewID(),
		loop:     loop,
		registry: registry,
		ctx:      pipeCtx,
		cancel:   cancel,
		ch:       make(chan blobPipeItem, 8),
	}
}

func (p *blobPipe) Push(chunk []byte) bool {
	if len(chunk) == 0 {
		return true
	}
	buf := append([]byte(nil), chunk...)
	select {
	case p.ch <- blobPipeItem{chunk: buf}:
		return true
	case <-p.ctx.Done():
		return false
	}
}

func (p *blobPipe) CloseWithError(err error) {
	p.once.Do(func() {
		p.registry.Remove(p.id)
		select {
		case p.ch <- blobPipeItem{done: true, err: err}:
		case <-p.ctx.Done():
		}
	})
}

func (p *blobPipe) Read(dst []byte) (int, error) {
	p.mu.Lock()
	if len(p.buf) > 0 {
		n := copy(dst, p.buf)
		p.buf = p.buf[n:]
		p.mu.Unlock()
		return n, nil
	}
	p.mu.Unlock()

	select {
	case item := <-p.ch:
		if item.done {
			if item.err != nil {
				return 0, item.err
			}
			return 0, io.EOF
		}
		n := copy(dst, item.chunk)
		if n < len(item.chunk) {
			p.mu.Lock()
			p.buf = append(p.buf, item.chunk[n:]...)
			p.mu.Unlock()
		}
		return n, nil
	case <-p.ctx.Done():
		return 0, p.ctx.Err()
	}
}

func (p *blobPipe) Close() error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	p.closed = true
	p.mu.Unlock()
	p.registry.Remove(p.id)
	p.cancel()
	p.cancelSourceReader()
	return nil
}

func (p *blobPipe) cancelSourceReader() {
	if p.loop == nil {
		return
	}
	// Invoke the JavaScript cancellation before returning from Close. The
	// source session may be released immediately after the HTTP handler closes
	// this reader, so merely queueing the callback could let the engine stop
	// before reader.cancel() is called.
	_, _ = runOnLoop(p.loop, func(runtime *goja.Runtime) (goja.Value, error) {
		fnVal := runtime.Get("__gopeed_blob_cancel_pipe_source")
		fn, ok := goja.AssertFunction(fnVal)
		if !ok {
			return goja.Undefined(), nil
		}
		_, err := fn(nil, runtime.ToValue(p.id), runtime.ToValue("blob request closed"))
		return goja.Undefined(), err
	})
}

func exportBytes(value goja.Value) []byte {
	if value == nil || goja.IsNull(value) || goja.IsUndefined(value) {
		return nil
	}
	if ab, ok := value.Export().(goja.ArrayBuffer); ok {
		return ab.Bytes()
	}
	if b, ok := value.Export().([]byte); ok {
		return b
	}
	return nil
}

type blobObjectURLReader struct {
	ctx  context.Context
	loop *eventloop.EventLoop
	id   string

	mu     sync.Mutex
	buf    []byte
	closed bool
}

func (r *blobObjectURLReader) Read(p []byte) (int, error) {
	r.mu.Lock()
	if len(r.buf) > 0 {
		n := copy(p, r.buf)
		r.buf = r.buf[n:]
		r.mu.Unlock()
		return n, nil
	}
	closed := r.closed
	r.mu.Unlock()
	if closed {
		return 0, io.EOF
	}
	if err := r.ctx.Err(); err != nil {
		return 0, err
	}
	chunkSize := len(p)
	if chunkSize <= 0 {
		chunkSize = 32 * 1024
	}
	value, err := runOnLoop(r.loop, func(runtime *goja.Runtime) (goja.Value, error) {
		fnVal := runtime.Get("__gopeed_blob_read_source")
		fn, ok := goja.AssertFunction(fnVal)
		if !ok {
			return nil, fmt.Errorf("blob read helper is not callable")
		}
		return fn(nil, runtime.ToValue(r.id), runtime.ToValue(chunkSize))
	})
	if err != nil {
		return 0, err
	}
	if goja.IsNull(value) || goja.IsUndefined(value) {
		return 0, io.EOF
	}
	var chunk []byte
	if ab, ok := value.Export().(goja.ArrayBuffer); ok {
		chunk = ab.Bytes()
	} else if b, ok := value.Export().([]byte); ok {
		chunk = b
	} else {
		return 0, fmt.Errorf("blob read helper returned %T", value.Export())
	}
	if len(chunk) == 0 {
		return 0, nil
	}
	n := copy(p, chunk)
	if n < len(chunk) {
		r.mu.Lock()
		r.buf = append(r.buf, chunk[n:]...)
		r.mu.Unlock()
	}
	return n, nil
}

func (r *blobObjectURLReader) Close() error {
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil
	}
	r.closed = true
	r.mu.Unlock()
	_, err := runOnLoop(r.loop, func(runtime *goja.Runtime) (goja.Value, error) {
		fnVal := runtime.Get("__gopeed_blob_close_source")
		fn, ok := goja.AssertFunction(fnVal)
		if !ok {
			return nil, nil
		}
		return fn(nil, runtime.ToValue(r.id))
	})
	return err
}

func runOnLoop(loop *eventloop.EventLoop, fn func(runtime *goja.Runtime) (goja.Value, error)) (goja.Value, error) {
	type result struct {
		value goja.Value
		err   error
	}
	ch := make(chan result, 1)
	ok := loop.RunOnLoop(func(runtime *goja.Runtime) {
		send := func(value goja.Value, err error) {
			select {
			case ch <- result{value: value, err: err}:
			default:
			}
		}
		defer func() {
			if r := recover(); r != nil {
				switch v := r.(type) {
				case error:
					send(nil, v)
				case goja.Value:
					send(nil, exportJSError(v))
				default:
					send(nil, fmt.Errorf("panic: %v", r))
				}
			}
		}()
		value, err := fn(runtime)
		if err != nil {
			send(nil, err)
			return
		}
		if promise, ok := value.Export().(*goja.Promise); ok {
			switch promise.State() {
			case goja.PromiseStateFulfilled:
				send(promise.Result(), nil)
				return
			case goja.PromiseStateRejected:
				send(nil, exportJSError(promise.Result()))
				return
			default:
				thenVal := value.ToObject(runtime).Get("then")
				thenFn, ok := goja.AssertFunction(thenVal)
				if !ok {
					send(nil, errors.New("promise.then is not callable"))
					return
				}
				onFulfilled := runtime.ToValue(func(call goja.FunctionCall) goja.Value {
					send(call.Argument(0), nil)
					return goja.Undefined()
				})
				onRejected := runtime.ToValue(func(call goja.FunctionCall) goja.Value {
					send(nil, exportJSError(call.Argument(0)))
					return goja.Undefined()
				})
				if _, err := thenFn(value, onFulfilled, onRejected); err != nil {
					send(nil, err)
				}
				return
			}
		}
		send(value, nil)
	})
	if !ok {
		return nil, errors.New("engine loop terminated")
	}
	res := <-ch
	return res.value, res.err
}

func exportJSError(value goja.Value) error {
	if value == nil || goja.IsUndefined(value) || goja.IsNull(value) {
		return errors.New("promise rejected")
	}
	if err, ok := value.Export().(error); ok {
		return err
	}
	stack := value.String()
	if ro, ok := value.(*goja.Object); ok {
		stackVal := ro.Get("stack")
		if stackVal != nil && stackVal.String() != "" {
			stack = stackVal.String()
		}
	}
	return errors.New(stack)
}