chrome.go

// package chrome provides helper functions to execute JS in a Chrome browser
//
// This would ordinarily be done via a Chrome struct but Go does not allow
// generic methods, only generic static functions, producing "method must have no type parameters".
package chrome

import (
	"context"
	"embed"
	"fmt"
	"io/fs"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/chromedp"
	"github.com/matrix-org/complement/ct"
)

//go:embed dist
var jsSDKDistDirectory embed.FS

// Void is a type which can be used when you want to run an async function without returning anything.
// It can stop large responses causing errors "Object reference chain is too long (-32000)"
// when we don't care about the response.
type Void *runtime.RemoteObject

// Run an anonymous async iffe in the browser. Set the type parameter to a basic data type
// which can be returned as JSON e.g string, map[string]any, []string. If you do not want
// to return anything, use chrome.Void. For example:
//
//	result, err := RunAsyncFn[string](t, ctx, "return await getSomeString()")
//	void, err := RunAsyncFn[chrome.Void](t, ctx, "doSomething(); await doSomethingElse();")
func RunAsyncFn[T any](t ct.TestLike, ctx context.Context, js string) (*T, error) {
	t.Helper()
	out := new(T)
	err := chromedp.Run(ctx,
		chromedp.Evaluate(`(async () => {`+js+`})()`, &out, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
			return p.WithAwaitPromise(true)
		}),
	)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// MustRunAsyncFn is RunAsyncFn but fails the test if an error is returned when executing.
//
// Run an anonymous async iffe in the browser. Set the type parameter to a basic data type
// which can be returned as JSON e.g string, map[string]any, []string. If you do not want
// to return anything, use chrome.Void
func MustRunAsyncFn[T any](t ct.TestLike, ctx context.Context, js string) *T {
	t.Helper()
	result, err := RunAsyncFn[T](t, ctx, js)
	if err != nil {
		ct.Fatalf(t, "MustRunAsyncFn: %s", err)
	}
	return result
}

type Browser struct {
	BaseURL string
	Ctx     context.Context
	Cancel  func()
}

// colouredLogWriter is an implementation of `io.Writer` which wraps an `os.File`, adding the date and logPrefix to each line,
// and decorating the text in a specific colour.
type coloredLogWriter struct {
	colour    string
	logPrefix string
	output    *os.File
}

func (w coloredLogWriter) Write(p []byte) (n int, err error) {
	_, err = w.output.WriteString(w.colour)
	if err != nil {
		return
	}

	last := byte('\n')
	for _, c := range p {
		// If this is the first byte after a newline, add the date and prefix
		if last == '\n' {
			_, err = w.output.WriteString(time.Now().Format(time.RFC3339) + " " + w.logPrefix + " ")
			if err != nil {
				return
			}
		}
		_, err = w.output.Write([]byte{c})
		if err != nil {
			return
		}
		n += 1
		last = c
	}

	ansiResetForeground := "\x1b[39m"
	_, err = w.output.WriteString(ansiResetForeground)
	return
}

func RunHeadless(logPrefix string, onConsoleLog func(s string), requiresPersistence bool, listenPort int) (*Browser, error) {
	ansiRedForeground := "\x1b[31m"
	ansiYellowForeground := "\x1b[33m"
	ansiBlueForeground := "\x1b[34m"

	// colorifyError returns a log format function which prints its input with a given prefix and colour.
	colorifyError := func(colour string, prefix string) func(format string, args ...any) {
		writer := coloredLogWriter{colour: colour, logPrefix: logPrefix + "[chromedp " + prefix + "]", output: os.Stdout}
		return func(format string, args ...any) {
			fmt.Fprintf(writer, format, args...)
		}
	}
	opts := chromedp.DefaultExecAllocatorOptions[:]
	if requiresPersistence {
		os.Mkdir("chromedp", os.ModePerm) // ignore errors to allow repeated runs
		wd, _ := os.Getwd()
		userDir := filepath.Join(wd, "chromedp")
		opts = append(opts,
			chromedp.UserDataDir(userDir),
		)
	}
	// increase the WS timeout from 20s (default) to 30s as we see timeouts with 20s in CI
	opts = append(opts, chromedp.WSURLReadTimeout(30*time.Second))

	// Capture stdout/stderr from Chrome, and log it.
	opts = append(opts, chromedp.CombinedOutput(coloredLogWriter{colour: ansiBlueForeground, logPrefix: logPrefix + " chrome:", output: os.Stdout}))

	// Hook into chromedp to log the command that is about to be executed. The easiest way to do that seems to be to
	// set a ModifyCmdFunc.
	opts = append(opts, chromedp.ModifyCmdFunc(func(cmd *exec.Cmd) {
		writer := coloredLogWriter{colour: ansiBlueForeground, logPrefix: logPrefix, output: os.Stdout}
		fmt.Fprintf(writer, "Executing: %v\n", cmd.Args)

		// Replicate the behaviour of the default ModifyCmdFunc: tell the kernel to send the child a SIGKILL when the
		// parent thread dies.
		if cmd.SysProcAttr == nil {
			cmd.SysProcAttr = new(syscall.SysProcAttr)
		}
		cmd.SysProcAttr.Pdeathsig = syscall.SIGKILL
	}))

	allocCtx, allocCancel := chromedp.NewExecAllocator(context.Background(), opts...)
	ctx, cancel := chromedp.NewContext(allocCtx, chromedp.WithBrowserOption(
		chromedp.WithBrowserLogf(colorifyError(ansiYellowForeground, "LOG")),
		chromedp.WithBrowserErrorf(colorifyError(ansiRedForeground, "ERROR")),
		// chromedp.WithBrowserDebugf(log.Printf),
	))

	// Make sure we can talk to Chrome before we start messing about with js-sdk
	fmt.Println("Starting chrome via chromedp")
	err := chromedp.Run(ctx, chromedp.Evaluate(`(() => 1)()`, nil))
	if err != nil {
		return nil, fmt.Errorf("failed to run test evaluation via chromedp: %w", err)
	}

	// Listen for console logs for debugging AND to communicate live updates
	chromedp.ListenTarget(ctx, func(ev interface{}) {
		switch ev := ev.(type) {
		case *runtime.EventConsoleAPICalled:
			for _, arg := range ev.Args {
				s, err := strconv.Unquote(string(arg.Value))
				if err != nil {
					s = string(arg.Value)
				}
				onConsoleLog(s)
			}
		}
	})

	// strip /dist so /index.html loads correctly as does /assets/xxx.js
	c, err := fs.Sub(jsSDKDistDirectory, "dist")
	if err != nil {
		return nil, fmt.Errorf("failed to strip /dist off JS SDK files: %s", err)
	}

	baseJSURL := ""
	// run js-sdk (need to run this as a web server to avoid CORS errors you'd otherwise get with file: URLs)
	var wg sync.WaitGroup
	wg.Add(1)
	mux := &http.ServeMux{}
	mux.Handle("/", http.FileServer(http.FS(c)))
	srv := &http.Server{
		Addr:    fmt.Sprintf("127.0.0.1:%d", listenPort),
		Handler: mux,
	}
	startServer := func() {
		ln, err := net.Listen("tcp", srv.Addr)
		if err != nil {
			panic(err)
		}
		baseJSURL = "http://" + ln.Addr().String()
		fmt.Println(logPrefix, "JS SDK wrapper listening on", baseJSURL)
		wg.Done()
		err = srv.Serve(ln)
		fmt.Printf("%s Closing webserver for JS SDK wrapper at %s: %v\n", logPrefix, baseJSURL, err)
	}
	go startServer()
	wg.Wait()

	// navigate to the page
	err = chromedp.Run(ctx,
		chromedp.Navigate(baseJSURL),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to navigate to %s: %s", baseJSURL, err)
	}

	return &Browser{
		Ctx: ctx,
		Cancel: func() {
			cancel()
			allocCancel()
			srv.Close()
		},
		BaseURL: baseJSURL,
	}, nil
}