diff --git a/pkg/modules/chromium/chromium.go b/pkg/modules/chromium/chromium.go index 29e0cf60..52bd82cf 100644 --- a/pkg/modules/chromium/chromium.go +++ b/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, &buffer)) + err := chromedp.Run(taskCtx, printToPDF(URL, options, outputPath)) 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 } diff --git a/pkg/modules/chromium/stream.go b/pkg/modules/chromium/stream.go new file mode 100644 index 00000000..70d20576 --- /dev/null +++ b/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) +)