fix(dialtesting): upstream Lightpanda 0.4 screenshot integration

dialtesting/browserdial/lightpanda/lightpanda.go

@@ -1,3 +1,8 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the MIT License.
+// This product includes software developed at Guance Cloud (https://www.guance.com/).
+// Copyright 2021-present Guance, Inc.
+
 package lightpanda
 
 import (
@@ -20,7 +25,9 @@ import (
 	"github.com/GuanceCloud/cliutils/dialtesting/browserdial/evidence"
 	"github.com/GuanceCloud/cliutils/dialtesting/browserdial/runner"
 	"github.com/GuanceCloud/cliutils/dialtesting/browserdial/util"
+	"github.com/chromedp/cdproto/emulation"
 	"github.com/chromedp/cdproto/network"
+	"github.com/chromedp/cdproto/page"
 	cdpruntime "github.com/chromedp/cdproto/runtime"
 	"github.com/chromedp/cdproto/security"
 	"github.com/chromedp/chromedp"
@@ -45,6 +52,8 @@ type Engine struct {
 	responses map[network.RequestID]struct{}
 }
 
+var _ runner.Screenshotter = (*Engine)(nil)
+
 type requestInfo struct {
 	URL          string
 	Method       string
@@ -91,7 +100,16 @@ func NewEngine(ctx context.Context, options runner.EngineOptions) (runner.Engine
 	}
 
 	chromedp.ListenTarget(tabCtx, engine.listen)
-	if err := engine.run(tabCtx, network.Enable(), cdpruntime.Enable()); err != nil {
+	actions := []chromedp.Action{network.Enable(), cdpruntime.Enable()}
+	if options.ViewportWidth > 0 && options.ViewportHeight > 0 {
+		actions = append(actions, emulation.SetDeviceMetricsOverride(
+			int64(options.ViewportWidth),
+			int64(options.ViewportHeight),
+			1,
+			false,
+		))
+	}
+	if err := engine.run(tabCtx, actions...); err != nil {
 		engine.cancel()
 		return nil, err
 	}
@@ -103,7 +121,13 @@ func lightpandaArguments(options runner.EngineOptions) ([]string, error) {
 }
 
 func lightpandaArgumentsWithSystemCADirectories(options runner.EngineOptions, systemDirectories []string) ([]string, error) {
-	arguments := []string{}
+	// Lightpanda 0.4.0 no longer loads iframes or workers by default. Keep
+	// browser dialtesting behavior compatible with the previously bundled
+	// engine while adopting the new explicit resource-loading interface.
+	arguments := []string{
+		"--load-resources", "iframe",
+		"--load-resources", "worker",
+	}
 	caCertFile := strings.TrimSpace(options.CACertFile)
 	caCertDir := strings.TrimSpace(options.CACertDir)
 	if caCertFile != "" || caCertDir != "" {
@@ -150,7 +174,9 @@ func lightpandaArgumentsWithSystemCADirectories(options runner.EngineOptions, sy
 		cidrs = append(cidrs, cidr)
 	}
 	if len(cidrs) > 0 {
-		arguments = append(arguments, "--block-cidrs", strings.Join(cidrs, ","))
+		for _, cidr := range cidrs {
+			arguments = append(arguments, "--block-cidrs", cidr)
+		}
 	} else if options.BlockPrivateNetwork {
 		arguments = append(arguments, "--block-private-networks")
 	}
@@ -283,6 +309,39 @@ func (e *Engine) Eval(ctx context.Context, expression string) (string, error) {
 	return util.JSONString(result, 8_000), err
 }
 
+func (e *Engine) CaptureScreenshot(ctx context.Context, path string, fullPage bool) (string, error) {
+	actionCtx, cancel := e.actionContext(ctx)
+	defer cancel()
+
+	// Lightpanda only supports PNG screenshots. Keep the returned path's
+	// extension consistent even if a caller requests a full-page capture.
+	if !strings.EqualFold(filepath.Ext(path), ".png") {
+		path = strings.TrimSuffix(path, filepath.Ext(path)) + ".png"
+	}
+	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
+		return "", err
+	}
+
+	var image []byte
+	capture := page.CaptureScreenshot().
+		WithFormat(page.CaptureScreenshotFormatPng).
+		WithCaptureBeyondViewport(fullPage)
+	if err := e.run(actionCtx, chromedp.ActionFunc(func(ctx context.Context) error {
+		var err error
+		image, err = capture.Do(ctx)
+		return err
+	})); err != nil {
+		return "", err
+	}
+	if len(image) == 0 {
+		return "", fmt.Errorf("lightpanda returned an empty screenshot")
+	}
+	if err := os.WriteFile(path, image, 0o600); err != nil {
+		return "", err
+	}
+	return path, nil
+}
+
 func (e *Engine) CaptureDOM(ctx context.Context) (evidence.DomSnapshot, error) {
 	snapshot := evidence.DomSnapshot{CapturedAt: util.NowISO()}
 	if currentURL, err := e.URL(ctx); err == nil {
@@ -337,9 +396,12 @@ func (e *Engine) evaluate(ctx context.Context, expression string, out any) error
 }
 
 func (e *Engine) actionContext(ctx context.Context) (context.Context, context.CancelFunc) {
-	actionCtx, cancel := context.WithCancel(e.ctx)
+	var actionCtx context.Context
+	var cancel context.CancelFunc
 	if deadline, ok := ctx.Deadline(); ok {
 		actionCtx, cancel = context.WithDeadline(e.ctx, deadline)
+	} else {
+		actionCtx, cancel = context.WithCancel(e.ctx)
 	}
 	go func() {
 		select {

dialtesting/browserdial/lightpanda/lightpanda_args_test.go

@@ -26,6 +26,7 @@ func TestLightpandaArguments(t *testing.T) {
 		t.Fatal(err)
 	}
 	want := []string{
+		"--load-resources", "iframe", "--load-resources", "worker",
 		"--ca-cert", caFile,
 		"--ca-path", caDir,
 		"--http-proxy", "http://user:password@proxy.example.com:8080",
@@ -38,14 +39,14 @@ func TestLightpandaArguments(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if !reflect.DeepEqual(arguments, []string{"--block-private-networks"}) {
+	if !reflect.DeepEqual(arguments, []string{"--load-resources", "iframe", "--load-resources", "worker", "--block-private-networks"}) {
 		t.Fatalf("private network block argument is missing: %#v", arguments)
 	}
 	arguments, err = lightpandaArguments(runner.EngineOptions{})
 	if err != nil {
 		t.Fatal(err)
 	}
-	if len(arguments) != 0 {
+	if !reflect.DeepEqual(arguments, []string{"--load-resources", "iframe", "--load-resources", "worker"}) {
 		t.Fatalf("zero-value options should preserve existing behavior: %#v", arguments)
 	}
 }
@@ -61,9 +62,10 @@ func TestLightpandaArgumentsPreserveSystemCAAndCustomCIDRs(t *testing.T) {
 		t.Fatal(err)
 	}
 	want := []string{
+		"--load-resources", "iframe", "--load-resources", "worker",
 		"--ca-path", systemDirectory,
 		"--ca-path", customDirectory,
-		"--block-cidrs", "10.0.0.0/8,192.168.0.0/16",
+		"--block-cidrs", "10.0.0.0/8", "--block-cidrs", "192.168.0.0/16",
 	}
 	if !reflect.DeepEqual(arguments, want) {
 		t.Fatalf("arguments = %#v, want %#v", arguments, want)

dialtesting/browserdial/lightpanda/lightpanda_screenshot_test.go

@@ -0,0 +1,121 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the MIT License.
+// This product includes software developed at Guance Cloud (https://www.guance.com/).
+// Copyright 2021-present Guance, Inc.
+
+package lightpanda
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/chromedp/cdproto/cdp"
+	"github.com/chromedp/cdproto/page"
+	"github.com/chromedp/chromedp"
+)
+
+type screenshotExecutor struct {
+	image  []byte
+	params page.CaptureScreenshotParams
+}
+
+func (e *screenshotExecutor) Execute(_ context.Context, method string, params, result any) error {
+	if method != page.CommandCaptureScreenshot {
+		return fmt.Errorf("unexpected CDP method %q", method)
+	}
+	captureParams, ok := params.(*page.CaptureScreenshotParams)
+	if !ok {
+		return fmt.Errorf("unexpected screenshot params type %T", params)
+	}
+	e.params = *captureParams
+	captureResult, ok := result.(*page.CaptureScreenshotReturns)
+	if !ok {
+		return fmt.Errorf("unexpected screenshot result type %T", result)
+	}
+	captureResult.Data = base64.StdEncoding.EncodeToString(e.image)
+	return nil
+}
+
+func runWithExecutor(executor cdp.Executor) func(context.Context, ...chromedp.Action) error {
+	return func(ctx context.Context, actions ...chromedp.Action) error {
+		ctx = cdp.WithExecutor(ctx, executor)
+		for _, action := range actions {
+			if err := action.Do(ctx); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
+func TestCaptureScreenshot(t *testing.T) {
+	tests := []struct {
+		name           string
+		requestedPath  string
+		fullPage       bool
+		wantPath       string
+		beyondViewport bool
+	}{
+		{
+			name:          "viewport PNG",
+			requestedPath: "step-1.png",
+			wantPath:      "step-1.png",
+		},
+		{
+			name:           "full page remains PNG",
+			requestedPath:  "step-2.jpg",
+			fullPage:       true,
+			wantPath:       "step-2.png",
+			beyondViewport: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			image := []byte("png screenshot evidence")
+			executor := &screenshotExecutor{image: image}
+			engine := &Engine{
+				ctx: context.Background(),
+				run: runWithExecutor(executor),
+			}
+			requestedPath := filepath.Join(t.TempDir(), "nested", test.requestedPath)
+			gotPath, err := engine.CaptureScreenshot(context.Background(), requestedPath, test.fullPage)
+			if err != nil {
+				t.Fatalf("CaptureScreenshot() error = %v", err)
+			}
+			if want := filepath.Join(filepath.Dir(requestedPath), test.wantPath); gotPath != want {
+				t.Fatalf("CaptureScreenshot() path = %q, want %q", gotPath, want)
+			}
+			gotImage, err := os.ReadFile(gotPath)
+			if err != nil {
+				t.Fatalf("read screenshot: %v", err)
+			}
+			if !bytes.Equal(gotImage, image) {
+				t.Fatalf("screenshot data = %q, want %q", gotImage, image)
+			}
+			if executor.params.Format != page.CaptureScreenshotFormatPng {
+				t.Fatalf("screenshot format = %q, want png", executor.params.Format)
+			}
+			if executor.params.CaptureBeyondViewport != test.beyondViewport {
+				t.Fatalf("captureBeyondViewport = %t, want %t", executor.params.CaptureBeyondViewport, test.beyondViewport)
+			}
+		})
+	}
+}
+
+func TestCaptureScreenshotRejectsEmptyImage(t *testing.T) {
+	executor := &screenshotExecutor{}
+	engine := &Engine{
+		ctx: context.Background(),
+		run: runWithExecutor(executor),
+	}
+	path := filepath.Join(t.TempDir(), "step.png")
+	if _, err := engine.CaptureScreenshot(context.Background(), path, false); err == nil {
+		t.Fatal("CaptureScreenshot() error = nil, want empty screenshot error")
+	}
+}