ffctl/lib/marionette/gecko.go

package marionette

import (
	"context"
	"errors"
	"io"
	"net"
	"os"
	"os/exec"
)

func haveProg(prog string) bool {
	_, err := exec.LookPath(prog)
	return err == nil
}

func WithGecko(ctx context.Context, fn func(context.Context, *TCPTransport) error) (_retErr error) {
	maybeSetErr := func(_err error) {
		if _retErr == nil && _err != nil {
			_retErr = _err
		}
	}

	profile, err := os.MkdirTemp("", "")
	if err != nil {
		return err
	}
	defer func() { maybeSetErr(os.RemoveAll(profile)) }()

	var prog string
	switch {
	case haveProg("firefox"):
		prog = "firefox"
	case haveProg("iceweasel"):
		prog = "iceweasel"
	default:
		return errors.New("no gecko browser found")
	}
	cmd := exec.CommandContext(ctx, prog,
		"--new-instance",
		"--profile", profile,
		"--marionette")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		return err
	}

	defer func() { maybeSetErr(cmd.Wait()) }()
	defer func() {
		if _retErr != nil {
			cmd.Process.Signal(os.Interrupt)
		}
	}()

	var conn *net.TCPConn
	for conn == nil {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			_conn, _ := net.Dial("tcp", "localhost:2828")
			conn, _ = _conn.(*net.TCPConn)
		}
	}
	defer func() { maybeSetErr(conn.Close()) }()

	trans, err := NewTransport(conn)
	if err != nil {
		return err
	}

	if _, err := trans.CmdNewSession(nil); err != nil {
		return err
	}

	defer func() {
		_, err := trans.CmdQuit()
		maybeSetErr(err)
		maybeSetErr(conn.CloseWrite())
		_, err = io.Copy(io.Discard, conn)
		maybeSetErr(err)
	}()

	return fn(ctx, trans)
}