kage/browser/pool.go

// Package browser drives a real headless Chrome through the DevTools Protocol so
// JavaScript-built pages are captured as they actually render. kage always goes
// through here: navigate, let the page settle, then serialise the final DOM —
// the same markup a human would have seen — which the rest of the pipeline then
// strips of scripts and localises.
package browser

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/launcher"
	"github.com/go-rod/rod/lib/proto"
	"github.com/go-rod/stealth"
)

// Options configure a Pool.
type Options struct {
	Headless      bool          // run Chrome without a window
	Workers       int           // max concurrent pages
	Settle        time.Duration // network-idle quiet period after load
	RenderTimeout time.Duration // hard cap per page render
	Scroll        bool          // auto-scroll to trigger lazy-loaded media
	ChromeBin     string        // explicit binary; empty = autodetect
	ControlURL    string        // attach to an existing Chrome instead of launching
}

// DefaultOptions returns the baseline render settings.
func DefaultOptions() Options {
	return Options{
		Headless:      true,
		Workers:       4,
		Settle:        1500 * time.Millisecond,
		RenderTimeout: 30 * time.Second,
	}
}

// Pool owns one Chrome process shared across a run and bounds the number of
// pages open at once.
type Pool struct {
	opts Options
	sem  chan struct{}

	mu      sync.Mutex
	browser *rod.Browser
	closed  bool
}

// New creates a Pool. Chrome is launched lazily on the first Render.
func New(opts Options) *Pool {
	if opts.Workers < 1 {
		opts.Workers = 1
	}
	return &Pool{opts: opts, sem: make(chan struct{}, opts.Workers)}
}

// RenderResult is the outcome of rendering one page.
type RenderResult struct {
	HTML     string // the serialised final DOM
	FinalURL string // URL after any client-side redirects
	Title    string
}

// Render navigates to rawURL, lets it settle, and returns the final rendered
// HTML. It acquires a page slot from the pool and releases it when done.
func (p *Pool) Render(ctx context.Context, rawURL string) (RenderResult, error) {
	select {
	case p.sem <- struct{}{}:
		defer func() { <-p.sem }()
	case <-ctx.Done():
		return RenderResult{}, ctx.Err()
	}

	b, err := p.getBrowser()
	if err != nil {
		return RenderResult{}, err
	}

	page, err := stealth.Page(b)
	if err != nil {
		return RenderResult{}, fmt.Errorf("new page: %w", err)
	}
	defer func() { _ = page.Close() }()

	page = page.Context(ctx).Timeout(p.opts.RenderTimeout)

	if err := page.Navigate(rawURL); err != nil {
		return RenderResult{}, fmt.Errorf("navigate %s: %w", rawURL, err)
	}
	if err := page.WaitLoad(); err != nil {
		return RenderResult{}, fmt.Errorf("wait load %s: %w", rawURL, err)
	}
	settle(page, p.opts.Settle)
	if p.opts.Scroll {
		autoScroll(page)
		settle(page, p.opts.Settle)
	}

	html, err := page.HTML()
	if err != nil {
		return RenderResult{}, fmt.Errorf("serialise %s: %w", rawURL, err)
	}

	res := RenderResult{HTML: html, FinalURL: rawURL}
	if info, err := page.Info(); err == nil && info != nil {
		res.FinalURL = info.URL
		res.Title = info.Title
	}
	return res, nil
}

// getBrowser lazily connects to or launches Chrome.
func (p *Pool) getBrowser() (*rod.Browser, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return nil, fmt.Errorf("pool is closed")
	}
	if p.browser != nil {
		return p.browser, nil
	}

	controlURL := p.opts.ControlURL
	if controlURL == "" {
		l := launcher.New().
			Headless(p.opts.Headless).
			Set("disable-blink-features", "AutomationControlled").
			Set("disable-gpu", "")

		// Chrome's sandbox is the main line of defense when rendering pages from
		// the open web, so kage keeps it on by default (issue #10). It is dropped
		// only where it genuinely cannot initialize: inside a container, or when
		// running as root, where Chrome otherwise refuses to start. The decision
		// is logged so it is never silent.
		if off, reason := disableSandbox(); off {
			l = l.Set("no-sandbox", "")
			warnSandboxDisabled(reason)
		}

		// In a container, the default /dev/shm is only 64 MB, too small for
		// Chrome's renderer on large pages, so steer it to a temp file instead.
		// Outside a container /dev/shm is roomy and faster, so leave it alone.
		if inContainer() {
			l = l.Set("disable-dev-shm-usage", "")
		}

		if bin := p.chromeBin(); bin != "" {
			l = l.Bin(bin)
		}
		u, err := l.Launch()
		if err != nil {
			return nil, fmt.Errorf("launch Chrome: %w", err)
		}
		controlURL = u
	}

	b := rod.New().ControlURL(controlURL)
	if err := b.Connect(); err != nil {
		return nil, fmt.Errorf("connect Chrome: %w", err)
	}
	p.browser = b
	return b, nil
}

// Close shuts down the managed Chrome process.
func (p *Pool) Close() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.closed = true
	if p.browser == nil {
		return nil
	}
	err := p.browser.Close()
	p.browser = nil
	return err
}

// LookChrome reports the path of a usable Chrome/Chromium binary and whether one
// was found, checking KAGE_CHROME, CHROME_BIN, rod's own lookup, and the common
// system install locations. Tests use it to skip when no browser is present.
func LookChrome() (string, bool) {
	for _, env := range []string{"KAGE_CHROME", "CHROME_BIN"} {
		if v := os.Getenv(env); v != "" {
			return v, true
		}
	}
	if bin, ok := launcher.LookPath(); ok {
		return bin, true
	}
	for _, c := range systemChromeCandidates() {
		if _, err := os.Stat(c); err == nil {
			return c, true
		}
	}
	return "", false
}

// chromeBin returns an explicit Chrome path from options or the environment, or
// "" to let the launcher find/download one.
func (p *Pool) chromeBin() string {
	if p.opts.ChromeBin != "" {
		return p.opts.ChromeBin
	}
	for _, env := range []string{"KAGE_CHROME", "CHROME_BIN"} {
		if v := os.Getenv(env); v != "" {
			return v
		}
	}
	for _, c := range systemChromeCandidates() {
		if _, err := os.Stat(c); err == nil {
			return c
		}
	}
	return ""
}

func systemChromeCandidates() []string {
	switch runtime.GOOS {
	case "darwin":
		return []string{
			"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
			"/Applications/Chromium.app/Contents/MacOS/Chromium",
		}
	case "windows":
		return []string{
			`C:\Program Files\Google\Chrome\Application\chrome.exe`,
			`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
		}
	default:
		return []string{
			"/usr/bin/google-chrome",
			"/usr/bin/google-chrome-stable",
			"/usr/bin/chromium",
			"/usr/bin/chromium-browser",
		}
	}
}

// disableSandbox decides whether Chrome should launch without its sandbox, with
// a short reason for the log. The secure default is to keep the sandbox on; it
// is dropped only where it cannot run: inside a container, or when running as
// root (Chrome refuses to start a sandbox as root).
func disableSandbox() (off bool, reason string) {
	if inContainer() {
		return true, "container"
	}
	if isRoot() {
		return true, "root"
	}
	return false, ""
}

// warnSandboxDisabled prints why the sandbox was turned off, so dropping a
// security boundary is always visible rather than silent.
func warnSandboxDisabled(reason string) {
	switch reason {
	case "container":
		fmt.Fprintln(os.Stderr, "kage: container detected, Chrome sandbox disabled")
	case "root":
		fmt.Fprintln(os.Stderr, "kage: running as root, Chrome sandbox disabled (run as a non-root user to keep it on)")
	}
}

// inContainer reports whether kage is running inside a container, where Chrome
// needs container-specific flags. It honors IN_DOCKER (set it in your image)
// and the /.dockerenv marker that Docker writes into every container.
//
// Keeping the sandbox on by default and dropping it only here was prompted by
// Dimitrios Prasakis (issue #10); the IN_DOCKER opt-in was suggested on Hacker
// News (https://news.ycombinator.com/item?id=48534865). Thanks to both.
func inContainer() bool {
	if envTrue("IN_DOCKER") {
		return true
	}
	if _, err := os.Stat("/.dockerenv"); err == nil {
		return true
	}
	return false
}

// isRoot reports whether the process runs as the superuser. On Windows
// os.Geteuid returns -1, so this is false there.
func isRoot() bool {
	return os.Geteuid() == 0
}

// envTrue reports whether the named environment variable is set to a truthy
// value.
func envTrue(name string) bool {
	v, ok := envBool(name)
	return ok && v
}

// envBool parses a boolean-ish environment variable. It returns ok=false when
// the variable is unset or empty. "1", "true", "yes", "on" are true and "0",
// "false", "no", "off" are false (case-insensitive); any other non-empty value
// counts as true, so IN_DOCKER=docker reads as set.
func envBool(name string) (val, ok bool) {
	s := strings.TrimSpace(os.Getenv(name))
	if s == "" {
		return false, false
	}
	switch strings.ToLower(s) {
	case "1", "true", "yes", "on":
		return true, true
	case "0", "false", "no", "off":
		return false, true
	default:
		return true, true
	}
}

// settle waits for the network to go quiet for d, recovering from any rod
// panic and capping the wait so a chatty page can never hang the worker.
func settle(page *rod.Page, d time.Duration) {
	if d <= 0 {
		return
	}
	defer func() { _ = recover() }()
	done := make(chan struct{})
	go func() {
		defer func() { _ = recover(); close(done) }()
		wait := page.WaitRequestIdle(d, nil, nil, []proto.NetworkResourceType{})
		wait()
	}()
	select {
	case <-done:
	case <-time.After(d + 5*time.Second):
	}
}

// autoScroll scrolls to the bottom in steps to trigger lazy-loaded images.
func autoScroll(page *rod.Page) {
	defer func() { _ = recover() }()
	_, _ = page.Eval(`() => new Promise((resolve) => {
		let total = 0;
		const step = 800;
		const timer = setInterval(() => {
			window.scrollBy(0, step);
			total += step;
			if (total >= document.body.scrollHeight) {
				clearInterval(timer);
				window.scrollTo(0, 0);
				resolve(true);
			}
		}, 100);
	})`)
}