Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ require (
github.com/projectdiscovery/retryablehttp-go v1.3.29
github.com/projectdiscovery/tlsx v1.4.0
github.com/projectdiscovery/useragent v0.0.109
github.com/projectdiscovery/utils v0.11.8
github.com/projectdiscovery/utils v0.11.9-0.20261010192532-28bb00e69775
github.com/projectdiscovery/wappalyzergo v0.3.4
github.com/rs/xid v1.6.0
github.com/spaolacci/murmur3 v1.1.0
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -296,8 +296,8 @@ github.com/projectdiscovery/tlsx v1.4.0 h1:OU2mNK6TOB/YiFX3jhmgwOepFWxPoj5rn8npW
github.com/projectdiscovery/tlsx v1.4.0/go.mod h1:WCrUXAHjdKbDNjsu5EbVBS7qvzap3fxi/E85v/Flyx8=
github.com/projectdiscovery/useragent v0.0.109 h1:b86BPdZvmgbJousa04jAugDRY4Y+rAvEamE0TFCPmeE=
github.com/projectdiscovery/useragent v0.0.109/go.mod h1:sRyBqDAcD31RKYrO8fyXGowLd36fNwY1g0MNW9mp7zk=
github.com/projectdiscovery/utils v0.11.8 h1:oBP3eRG/S/Pdid6CkL58IVyq+KZdWdXohDtEKIKvWag=
github.com/projectdiscovery/utils v0.11.8/go.mod h1:AcQ1n3ETcYj4cHiPXsSaAyFnvN4Bn8UgN5tCmh+7KWM=
github.com/projectdiscovery/utils v0.11.9-0.20261010192532-28bb00e69775 h1:hnFyxKRnCZoB7C4kjsF8GqmYVoYR32X09Zu3gBTxnQQ=
github.com/projectdiscovery/utils v0.11.9-0.20261010192532-28bb00e69775/go.mod h1:d3BcbsHhuxW9T3jpvrgzwshvqhaL9qPSkoi+3h/DY8E=
github.com/projectdiscovery/wappalyzergo v0.3.4 h1:yO+xel+C0d/vOFWcUKVZRpIDAfcEmsjpVDjQ0jifAUw=
github.com/projectdiscovery/wappalyzergo v0.3.4/go.mod h1:6xTRCUQhX8m+L1eqJDQIALJlRv2Bnhrmku80YU9bBXE=
github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEvV+S9iJ2IdQo=
Expand Down
88 changes: 78 additions & 10 deletions runner/headless.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package runner

import (
"context"
"fmt"
"os"
"strings"
Expand All @@ -11,6 +12,7 @@ import (
"github.com/go-rod/rod/lib/launcher/flags"
"github.com/go-rod/rod/lib/proto"
"github.com/pkg/errors"
"github.com/projectdiscovery/utils/chromeshell"
fileutil "github.com/projectdiscovery/utils/file"
mapsutil "github.com/projectdiscovery/utils/maps"
osutils "github.com/projectdiscovery/utils/os"
Expand Down Expand Up @@ -61,6 +63,28 @@ func NewBrowser(proxy string, useLocal bool, optionalArgs map[string]string) (*B
Set("window-size", fmt.Sprintf("%d,%d", 1080, 1920)).
Set("mute-audio", "true").
Set("incognito", "true").
// Performance: keep background/occluded tabs running at full speed so
// many concurrent screenshot pages don't get render-throttled.
Set("disable-background-timer-throttling", "true").
Set("disable-backgrounding-occluded-windows", "true").
Set("disable-renderer-backgrounding", "true").
Set("disable-ipc-flooding-protection", "true").
Set("disable-hang-monitor", "true").
// Performance: strip background chrome services we never use.
Set("disable-background-networking", "true").
Set("disable-client-side-phishing-detection", "true").
Set("disable-component-update", "true").
Set("disable-default-apps", "true").
Set("disable-domain-reliability", "true").
Set("disable-extensions", "true").
Set("disable-sync", "true").
Set("no-first-run", "true").
Set("no-default-browser-check", "true").
Set("metrics-recording-only", "true").
Set("safebrowsing-disable-auto-update", "true").
// site-per-process stays enabled: screenshots render untrusted pages, so
// Chromium's cross-origin renderer isolation is worth the extra processes.
Set("disable-features", "Translate,BackForwardCache,AcceptCHFrame,MediaRouter,OptimizationHints").
Delete("use-mock-keychain").
Headless(true).
UserDataDir(dataStore)
Expand All @@ -82,6 +106,12 @@ func NewBrowser(proxy string, useLocal bool, optionalArgs map[string]string) (*B
} else {
return nil, errors.New("the chrome browser is not installed")
}
} else if chromeshell.Supported() {
// Prefer chrome-headless-shell on linux/amd64: smaller download and
// faster headless screenshots than full Chromium snapshots.
if shellPath, err := ensureChromeShell(); err == nil {
chromeLauncher.Bin(shellPath)
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

if proxy == "" {
Expand Down Expand Up @@ -119,13 +149,13 @@ func NewBrowser(proxy string, useLocal bool, optionalArgs map[string]string) (*B
}

func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle time.Duration, headers []string, fullPage bool, jsCodes []string) ([]byte, string, []NetworkRequest, error) {
page, networkRequests, err := b.setupPageAndNavigate(url, timeout, headers, jsCodes)
page, networkRequests, err := b.setupPageAndNavigate(url, timeout, idle, headers, jsCodes)
if err != nil {
return nil, "", []NetworkRequest{}, err
}
defer b.closePage(page)

screenshot, body, err := b.takeScreenshotAndGetBody(page, idle, fullPage)
screenshot, body, err := b.takeScreenshotAndGetBody(page, fullPage)
Comment thread
Mzack9999 marked this conversation as resolved.
if err != nil {
return nil, "", networkRequests, err
}
Expand All @@ -134,7 +164,7 @@ func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle tim
}

// setupPageAndNavigate opens a page, performs all adaptive actions including JS injection
func (b *Browser) setupPageAndNavigate(url string, timeout time.Duration, headers []string, jsCodes []string) (*rod.Page, []NetworkRequest, error) {
func (b *Browser) setupPageAndNavigate(url string, timeout time.Duration, idle time.Duration, headers []string, jsCodes []string) (*rod.Page, []NetworkRequest, error) {
page, err := b.engine.Page(proto.TargetCreateTarget{})
if err != nil {
return nil, []NetworkRequest{}, err
Expand Down Expand Up @@ -208,6 +238,13 @@ func (b *Browser) setupPageAndNavigate(url string, timeout time.Duration, header
}

page = page.Timeout(timeout)
var waitReqIdle func()
if idle > 0 {
// Register before Navigate so the first SPA XHRs are tracked by the
// request-idle waiter. Zero/negative idle skips this (WaitRequestIdle
// and WaitDOMStable reject non-positive durations).
waitReqIdle = page.WaitRequestIdle(idle, nil, nil, nil)
}

if err := page.Navigate(url); err != nil {
return page, networkRequests.Slice, err
Expand All @@ -220,19 +257,50 @@ func (b *Browser) setupPageAndNavigate(url string, timeout time.Duration, header
}
}

page.Timeout(5 * time.Second).WaitNavigation(proto.PageLifecycleEventNameFirstMeaningfulPaint)()
b.waitPageReady(page, idle, waitReqIdle)

return page, networkRequests.Slice, nil
}

// takeScreenshotAndGetBody performs the screenshot actions
func (b *Browser) takeScreenshotAndGetBody(page *rod.Page, idle time.Duration, fullPage bool) ([]byte, string, error) {
if err := page.WaitLoad(); err != nil {
return nil, "", err
// waitPageReady blocks until the page is visually settled so we don't capture a
// half-rendered SPA. It gates on three independent signals, each bounded by the
// page timeout:
// - window.onload
// - network request-idle (a quiet network window)
// - DOM stability (the rendered tree stops mutating)
//
// SPAs paint late: the network can briefly go quiet before hydration starts, so
// request-idle alone can fire on the boot screen. Requiring DOM stability on top
// closes that gap. DOM stability is re-checked after the network settles so the
// first snapshot is taken post-hydration rather than on the boot screen.
func (b *Browser) waitPageReady(page *rod.Page, idle time.Duration, waitReqIdle func()) {
_ = page.WaitLoad()
if idle <= 0 || waitReqIdle == nil {
return
}
_ = page.WaitIdle(idle)
waitReqIdle()
_ = page.WaitDOMStable(idle, 0)
}
Comment thread
Mzack9999 marked this conversation as resolved.

const chromeShellEnsureTimeout = 2 * time.Minute

// ensureChromeShell fetches the browser under a deadline. The context has to
// reach the download itself: racing a plain Ensure against a timer would only
// abandon the wait, leaving the transfer running and the shared download lock
// held, so the next caller would block behind it.
func ensureChromeShell() (string, error) {
ctx, cancel := context.WithTimeout(context.Background(), chromeShellEnsureTimeout)
defer cancel()
return chromeshell.EnsureContext(ctx)
}

// takeScreenshotAndGetBody performs the screenshot actions
func (b *Browser) takeScreenshotAndGetBody(page *rod.Page, fullPage bool) ([]byte, string, error) {
_ = page.WaitRepaint()

screenshot, err := page.Screenshot(fullPage, &proto.PageCaptureScreenshot{})
screenshot, err := page.Screenshot(fullPage, &proto.PageCaptureScreenshot{
OptimizeForSpeed: true,
})
if err != nil {
return nil, "", err
}
Expand Down
139 changes: 139 additions & 0 deletions runner/headless_screenshot_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
package runner

import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
)

const spaHydratedMarker = "hydrated-ok"

// spaHTML is a tiny SPA-style page: it paints a boot screen first, then
// fetches /app.js which hydrates #root after a short delay. Early captures
// only contain "boot"; a correct wait must include spaHydratedMarker.
const spaHTML = `<!doctype html>
<html><head><title>spa-bench</title></head>
<body>
<div id="root">boot</div>
<script>fetch("/app.js").then(function(r){return r.text()}).then(eval);</script>
</body></html>`

const spaJS = `
setTimeout(function(){
document.getElementById("root").textContent = "` + spaHydratedMarker + `";
}, 80);
`

func startSPAScreenshotServer(t testing.TB) *httptest.Server {
t.Helper()
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html; charset=utf-8")
_, _ = w.Write([]byte(spaHTML))
})
mux.HandleFunc("/app.js", func(w http.ResponseWriter, r *http.Request) {
time.Sleep(50 * time.Millisecond)
w.Header().Set("Content-Type", "application/javascript")
_, _ = w.Write([]byte(spaJS))
})
srv := httptest.NewServer(mux)
t.Cleanup(srv.Close)
return srv
}

func newScreenshotBrowser(t testing.TB, useLocal bool) *Browser {
t.Helper()
browser, err := NewBrowser("", useLocal, nil)
if err != nil {
t.Skipf("headless browser unavailable: %v", err)
}
t.Cleanup(browser.Close)
return browser
}

func captureSPA(t testing.TB, browser *Browser, url string) (screenshot []byte, body string) {
t.Helper()
screenshot, body, _, err := browser.ScreenshotWithBody(url, 15*time.Second, 200*time.Millisecond, nil, false, nil)
if err != nil {
t.Fatalf("screenshot: %v", err)
}
return screenshot, body
}

func TestScreenshotSPAWaitsForHydration(t *testing.T) {
srv := startSPAScreenshotServer(t)
browser := newScreenshotBrowser(t, false)

screenshot, body := captureSPA(t, browser, srv.URL)
if len(screenshot) == 0 {
t.Fatal("empty screenshot")
}
if !strings.Contains(body, spaHydratedMarker) {
t.Fatalf("early capture: body missing %q\n%s", spaHydratedMarker, body)
}
}

func TestScreenshotSPAConcurrentNoEarlyCapture(t *testing.T) {
const workers = 8
srv := startSPAScreenshotServer(t)
browser := newScreenshotBrowser(t, false)

var wg sync.WaitGroup
errs := make(chan error, workers)
for i := 0; i < workers; i++ {
wg.Add(1)
go func() {
defer wg.Done()
screenshot, body, _, err := browser.ScreenshotWithBody(srv.URL, 15*time.Second, 200*time.Millisecond, nil, false, nil)
if err != nil {
errs <- err
return
}
if len(screenshot) == 0 {
errs <- fmt.Errorf("empty screenshot")
return
}
if !strings.Contains(body, spaHydratedMarker) {
errs <- fmt.Errorf("early capture: body missing %q", spaHydratedMarker)
}
}()
}
wg.Wait()
close(errs)
for err := range errs {
t.Error(err)
}
}

// BenchmarkScreenshotSPA times concurrent captures of the local SPA.
//
// On linux/amd64, /default uses chrome-headless-shell when available and
// /local forces a system Chrome via NewBrowser(..., useLocal=true). Compare:
//
// go test -bench=BenchmarkScreenshotSPA -benchtime=20s -count=5 ./runner
func BenchmarkScreenshotSPA(b *testing.B) {
for _, useLocal := range []bool{false, true} {
name := "default"
if useLocal {
name = "local"
}
b.Run(name, func(b *testing.B) {
srv := startSPAScreenshotServer(b)
browser := newScreenshotBrowser(b, useLocal)
b.ResetTimer()
for i := 0; i < b.N; i++ {
screenshot, body, _, err := browser.ScreenshotWithBody(srv.URL, 15*time.Second, 200*time.Millisecond, nil, false, nil)
if err != nil {
b.Fatalf("screenshot: %v", err)
}
if len(screenshot) == 0 || !strings.Contains(body, spaHydratedMarker) {
b.Fatal("early or empty capture")
}
}
Comment thread
ehsandeep marked this conversation as resolved.
})
}
}
Loading