feat(chromium): add ReturnAsStream support

pkg/modules/chromium/chromium.go

@@ -1,6 +1,7 @@
 package chromium
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
@@ -425,7 +426,7 @@ func (mod Chromium) PDF(ctx context.Context, logger *zap.Logger, URL, outputPath
 		consoleExceptionsMu sync.RWMutex
 	)
 
-	printToPDF := func(URL string, options Options, result *[]byte) chromedp.Tasks {
+	printToPDF := func(URL string, options Options, outputPath string) chromedp.Tasks {
 		// We validate the underlying requests against our allow / deny lists.
 		// If a request does not pass the validation, we make it fail.
 		listenForEventRequestPaused(taskCtx, logger, mod.allowList, mod.denyList)
@@ -731,6 +732,7 @@ func (mod Chromium) PDF(ctx context.Context, logger *zap.Logger, URL, outputPath
 			}),
 			chromedp.ActionFunc(func(ctx context.Context) error {
 				printToPDF := page.PrintToPDF().
+					WithTransferMode(page.PrintToPDFTransferModeReturnAsStream).
 					WithLandscape(options.Landscape).
 					WithPrintBackground(options.PrintBackground).
 					WithScale(options.Scale).
@@ -748,12 +750,44 @@ func (mod Chromium) PDF(ctx context.Context, logger *zap.Logger, URL, outputPath
 
 				logger.Debug(fmt.Sprintf("print to PDF with: %+v", printToPDF))
 
-				data, _, err := printToPDF.Do(ctx)
+				_, stream, err := printToPDF.Do(ctx)
 				if err != nil {
 					return fmt.Errorf("print to PDF: %w", err)
 				}
 
-				*result = data
+				reader := &streamReader{
+					ctx:    ctx,
+					handle: stream,
+					r:      nil,
+					pos:    0,
+					eof:    false,
+				}
+
+				defer func() {
+					err := reader.Close()
+					if err != nil {
+						logger.Error(fmt.Sprintf("close reader: %s", err))
+					}
+				}()
+
+				file, err := os.OpenFile(outputPath, os.O_CREATE|os.O_WRONLY, 0600)
+				if err != nil {
+					return fmt.Errorf("open output path: %w", err)
+				}
+
+				defer func() {
+					err := file.Close()
+					if err != nil {
+						logger.Error(fmt.Sprintf("close output path: %s", err))
+					}
+				}()
+
+				buffer := bufio.NewReader(reader)
+
+				_, err = buffer.WriteTo(file)
+				if err != nil {
+					return fmt.Errorf("write result to output path: %w", err)
+				}
 
 				return nil
 			}),
@@ -764,8 +798,7 @@ func (mod Chromium) PDF(ctx context.Context, logger *zap.Logger, URL, outputPath
 	activeInstancesCount += 1
 	activeInstancesCountMu.Unlock()
 
-	var buffer []byte
+	err := chromedp.Run(taskCtx, printToPDF(URL, options, outputPath))
-	err := chromedp.Run(taskCtx, printToPDF(URL, options, &buffer))
 
 	activeInstancesCountMu.Lock()
 	activeInstancesCount -= 1
@@ -807,11 +840,6 @@ func (mod Chromium) PDF(ctx context.Context, logger *zap.Logger, URL, outputPath
 		return fmt.Errorf("%v: %w", consoleExceptions, ErrConsoleExceptions)
 	}
 
-	err = os.WriteFile(outputPath, buffer, 0600)
-	if err != nil {
-		return fmt.Errorf("write result to output path: %w", err)
-	}
-
 	return nil
 }
 

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