fix(headless): avoid waiting for an already-fired load event

pkg/protocols/headless/engine/page_actions.go

@@ -732,6 +732,18 @@ func (p *Page) SelectInputElement(act *Action, out ActionData) error {
 
 // WaitPageLifecycleEvent waits for specified page lifecycle event name
 func (p *Page) WaitPageLifecycleEvent(act *Action, out ActionData, event proto.PageLifecycleEventName) error {
+	if event == proto.PageLifecycleEventNameLoad {
+		dur, err := getTimeout(p, act)
+		if err != nil {
+			return err
+		}
+		if err := p.page.Timeout(dur).WaitLoad(); err != nil {
+			return err
+		}
+		p.updateLastNavigatedURL()
+		return nil
+	}
+
 	fn, err := getNavigationFunc(p, act, event)
 	if err != nil {
 		return err

pkg/protocols/headless/engine/page_actions_test.go

@@ -18,6 +18,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-rod/rod/lib/proto"
 	"github.com/stretchr/testify/require"
 
 	"github.com/projectdiscovery/nuclei/v3/internal/tests/testheadless"
@@ -642,6 +643,47 @@ func TestActionWaitLoad(t *testing.T) {
 	})
 }
 
+func TestActionWaitLoadAfterPageAlreadyLoaded(t *testing.T) {
+	actions := []*Action{{ActionType: ActionTypeHolder{ActionType: ActionNavigate}, Data: map[string]string{"url": "{{BaseURL}}"}}}
+	testHeadlessSimpleResponse(t, "<html><body>ready</body></html>", actions, 10*time.Second, func(page *Page, err error, _ ActionData) {
+		require.NoError(t, err)
+		require.NoError(t, page.Page().WaitLoad())
+
+		wait := &Action{ActionType: ActionTypeHolder{ActionType: ActionWaitLoad}, Data: map[string]string{"timeout": "2s"}}
+		started := time.Now()
+		require.NoError(t, page.WaitPageLifecycleEvent(wait, nil, proto.PageLifecycleEventNameLoad))
+		require.Less(t, time.Since(started), 1500*time.Millisecond, "a completed load must not consume the full wait timeout")
+	})
+}
+
+func TestActionWaitLoadWaitsForPendingResource(t *testing.T) {
+	resourceDone := make(chan struct{})
+	actions := []*Action{
+		{ActionType: ActionTypeHolder{ActionType: ActionNavigate}, Data: map[string]string{"url": "{{BaseURL}}"}},
+		{ActionType: ActionTypeHolder{ActionType: ActionWaitLoad}},
+	}
+	testHeadless(t, actions, 10*time.Second, func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/slow.js" {
+			w.Header().Set("Content-Type", "text/javascript")
+			time.Sleep(300 * time.Millisecond)
+			_, _ = fmt.Fprint(w, "window.resourceLoaded = true")
+			close(resourceDone)
+			return
+		}
+		_, _ = fmt.Fprint(w, `<html><script src="/slow.js"></script></html>`)
+	}, func(page *Page, err error, _ ActionData) {
+		require.NoError(t, err)
+		select {
+		case <-resourceDone:
+		default:
+			t.Fatal("waitload returned before the pending resource completed")
+		}
+		loaded, evalErr := page.Page().Eval(`() => window.resourceLoaded`)
+		require.NoError(t, evalErr)
+		require.True(t, loaded.Value.Bool())
+	})
+}
+
 func TestActionGetResource(t *testing.T) {
 	response := `
 		<html>