feat(chromium): add ReturnAsStream support

pkg/modules/chromium/stream.go

@@ -0,0 +1,117 @@
+package chromium
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/chromedp/cdproto/cdp"
+	cdprotoio "github.com/chromedp/cdproto/io"
+)
+
+// Credits: https://raw.githubusercontent.com/mafredri/cdp/3c5eab7ffc5cbee667b0a813ce470ac423792811/protocol/io/stream_reader.go.
+type streamReader struct {
+	ctx    context.Context
+	handle cdprotoio.StreamHandle
+	r      io.Reader
+	pos    int
+	eof    bool
+}
+
+// Read a chunk of the stream.
+func (reader *streamReader) Read(p []byte) (n int, err error) {
+	if reader.r != nil {
+		// Continue reading from buffer.
+		return reader.read(p)
+	}
+
+	if reader.eof {
+		return 0, io.EOF
+	}
+
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	// Chromium might have an off-by-one when deciding the maximum size (at
+	// least for base64 encoded data), usually it will overflow. We subtract
+	// one to make sure it fits into p.
+	size := len(p) - 1
+	if size < 1 {
+		// Safety-check to avoid crashing Chrome (e.g. via SetSize(-1)).
+		size = 1
+	}
+
+	reply, err := reader.next(reader.pos, size)
+	if err != nil {
+		return 0, err
+	}
+
+	reader.eof = reply.EOF
+
+	switch {
+	case reply.Base64encoded:
+		b := []byte(reply.Data)
+		size := base64.StdEncoding.DecodedLen(len(b))
+
+		// Safety-check for fast-path to avoid panics.
+		if len(p) >= size {
+			n, err = base64.StdEncoding.Decode(p, b)
+			reader.pos += n
+
+			return n, err
+		}
+
+		reader.r = base64.NewDecoder(base64.StdEncoding, bytes.NewReader(b))
+	default:
+		reader.r = strings.NewReader(reply.Data)
+	}
+
+	return reader.read(p)
+}
+
+// Close closes the stream, discard any temporary backing storage.
+func (reader *streamReader) Close() error {
+	err := cdprotoio.Close(reader.handle).Do(reader.ctx)
+	if err == nil {
+		return nil
+	}
+
+	return fmt.Errorf("close Chromium stream: %w", err)
+}
+
+func (reader *streamReader) next(pos, size int) (cdprotoio.ReadReturns, error) {
+	params := cdprotoio.
+		Read(reader.handle).
+		WithOffset(int64(pos)).
+		WithSize(int64(size))
+
+	var res cdprotoio.ReadReturns
+	err := cdp.Execute(reader.ctx, cdprotoio.CommandRead, params, &res)
+
+	if err == nil {
+		return res, nil
+	}
+
+	return res, fmt.Errorf("execute IO.read command: %w", err)
+}
+
+func (reader *streamReader) read(p []byte) (n int, err error) {
+	n, err = reader.r.Read(p)
+	reader.pos += n
+
+	if !reader.eof && err == io.EOF {
+		reader.r = nil
+		err = nil
+	}
+
+	return n, err
+}
+
+// Interface guards.
+var (
+	_ io.Reader = (*streamReader)(nil)
+)