add ability to set chrome profile when launching

chrome/chrome.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"log"
 	"net/http"
 	"net/url"
 	"os"
@@ -19,7 +18,6 @@ var MaxTimeToStart = 10 * time.Second
 
 type Chrome struct {
 	WebSocketUrl string
-	UserDataDir  string
 	StartArgs    string
 	cmd          *exec.Cmd
 }
@@ -61,13 +59,7 @@ func (c Chrome) NewTab(cli *http.Client, address string) (target Target, err err
 	return
 }
 
-func (c Chrome) WaitCloseGracefully() error {
-	defer func() {
-		err := os.RemoveAll(c.UserDataDir)
-		if err != nil {
-			log.Println(err)
-		}
-	}()
+func (c Chrome) Wait() error {
 	return c.cmd.Wait()
 }
 
@@ -92,21 +84,15 @@ func bin() string {
 }
 
 func Launch(ctx context.Context, userFlags ...string) (value Chrome, err error) {
-	if value.UserDataDir, err = os.MkdirTemp("", "chrome-control-*"); err != nil {
-		return value, err
-	}
 	// https://github.com/GoogleChrome/chrome-launcher/blob/master/docs/chrome-flags-for-tools.md
 	// https://docs.google.com/spreadsheets/d/1n-vw_PCPS45jX3Jt9jQaAhFqBY6Ge1vWF_Pa0k7dCk4/edit#gid=1265672696
-	flags := []string{
-		"--remote-debugging-port=0",
-		"--user-data-dir=" + value.UserDataDir,
+	var flags = []string{"--remote-debugging-port=0"}
+	if os.Getuid() == 0 {
+		flags = append(flags, "--no-sandbox", "--disable-setuid-sandbox")
 	}
 	if len(userFlags) > 0 {
 		flags = append(flags, userFlags...)
 	}
-	if os.Getuid() == 0 {
-		flags = append(flags, "--no-sandbox", "--disable-setuid-sandbox")
-	}
 	binary := bin()
 	value.StartArgs = fmt.Sprint(binary, strings.Join(flags, " "))
 	value.cmd = exec.CommandContext(ctx, binary, flags...)

main.go

@@ -6,6 +6,8 @@ import (
 	"errors"
 	"log/slog"
 	"net/http"
+	"os"
+	"strings"
 
 	"github.com/ecwid/control/cdp"
 	"github.com/ecwid/control/chrome"
@@ -16,33 +18,63 @@ func Take(args ...string) (session *Session, cancel func() error, err error) {
 	return TakeWithContext(context.TODO(), nil, args...)
 }
 
-func TakeWithContext(ctx context.Context, logger *slog.Logger, chromeArgs ...string) (session *Session, cancel func() error, err error) {
-	browser, err := chrome.Launch(ctx, chromeArgs...)
+func hasProfile(args ...string) bool {
+	for _, a := range args {
+		a = strings.TrimSpace(a)
+		if strings.HasPrefix(a, "--profile-directory") || strings.HasPrefix(a, "--user-data-dir") {
+			return true
+		}
+	}
+	return false
+}
+
+func TakeWithContext(ctx context.Context, logger *slog.Logger, args ...string) (session *Session, cancel func() error, err error) {
+	var userDataDir string
+
+	if !hasProfile(args...) {
+		// If there is no profile or user-data-dir in the startup arguments, then user-data-dir must be set
+		userDataDir, err = os.MkdirTemp("", "chrome-control-*")
+		if err != nil {
+			return nil, nil, errors.Join(err, errors.New("can't create temporary user data dir"))
+		}
+		args = append(args, "--user-data-dir="+userDataDir)
+	}
+
+	browser, err := chrome.Launch(ctx, args...)
 	if err != nil {
 		return nil, nil, errors.Join(err, errors.New("chrome launch failed"))
 	}
+
 	tab, err := browser.NewTab(http.DefaultClient, "")
 	if err != nil {
 		return nil, nil, errors.Join(err, errors.New("failed to open a new tab"))
 	}
+
 	transport, err := cdp.DefaultDial(ctx, browser.WebSocketUrl, logger)
 	if err != nil {
 		return nil, nil, errors.Join(err, errors.New("websocket dial failed"))
 	}
+
 	session, err = NewSession(transport, target.TargetID(tab.ID))
 	if err != nil {
 		return nil, nil, errors.Join(err, errors.New("failed to create a new session"))
 	}
-	teardown := func() error {
+
+	cleanup := func() error {
 		if err := transport.Close(); err != nil {
 			return errors.Join(err, errors.New("can't close transport"))
 		}
-		if err = browser.WaitCloseGracefully(); err != nil {
+		if err = browser.Wait(); err != nil {
 			return errors.Join(err, errors.New("can't close browser gracefully"))
 		}
+		if userDataDir != "" {
+			if err := os.RemoveAll(userDataDir); err != nil {
+				return errors.Join(err, errors.New("can't clear user data dir"))
+			}
+		}
 		return nil
 	}
-	return session, teardown, nil
+	return session, cleanup, nil
 }
 
 func Subscribe[T any](s *Session, method string, filter func(T) bool) cdp.Future[T] {