headless: fix panic + refactor waitevent action (#4465)

* fix waitEvent action * avoid future panics * integration test + bug fix * headless: add max-duration support in waitevent * fix comment + max-duration input
projectdiscovery · Dec 6, 2023 · 47e7503 · 47e7503
1 parent 7da3921
commit 47e7503
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 31 deletions.
diff --git a/cmd/integration-test/headless.go b/cmd/integration-test/headless.go
@@ -12,6 +12,7 @@ import (
 
 var headlessTestcases = []TestCaseInfo{
 	{Path: "protocols/headless/headless-basic.yaml", TestCase: &headlessBasic{}},
+	{Path: "protocols/headless/headless-waitevent.yaml", TestCase: &headlessBasic{}},
 	{Path: "protocols/headless/headless-self-contained.yaml", TestCase: &headlessSelfContained{}},
 	{Path: "protocols/headless/headless-header-action.yaml", TestCase: &headlessHeaderActions{}},
 	{Path: "protocols/headless/headless-extract-values.yaml", TestCase: &headlessExtractValues{}},

diff --git a/integration_tests/protocols/headless/headless-waitevent.yaml b/integration_tests/protocols/headless/headless-waitevent.yaml
@@ -0,0 +1,24 @@
+id: headless-waitevent
+
+info:
+    name: WaitEvent
+    severity: info
+    author: pdteam
+
+headless:
+  - steps:
+      # note waitevent must be used before navigating to any page
+      # unlike waitload
+      - action: waitevent
+        args:
+            event: 'Page.loadEventFired'
+            max-duration: 15s
+
+      - action: navigate  
+        args:
+          url: "{{BaseURL}}/"
+
+    matchers:
+      - type: word
+        words:
+          - "<html>"
diff --git a/pkg/protocols/headless/engine/page_actions.go b/pkg/protocols/headless/engine/page_actions.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
@@ -21,6 +22,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/utils/vardump"
 	protocolutils "github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils"
 	httputil "github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils/http"
+	contextutil "github.com/projectdiscovery/utils/context"
 	errorutil "github.com/projectdiscovery/utils/errors"
 	fileutil "github.com/projectdiscovery/utils/file"
 	folderutil "github.com/projectdiscovery/utils/folder"
@@ -41,13 +43,33 @@ const (
 )
 
 // ExecuteActions executes a list of actions on a page.
-func (p *Page) ExecuteActions(input *contextargs.Context, actions []*Action, variables map[string]interface{}) (map[string]string, error) {
-	outData := make(map[string]string)
-	var err error
+func (p *Page) ExecuteActions(input *contextargs.Context, actions []*Action, variables map[string]interface{}) (outData map[string]string, err error) {
+	outData = make(map[string]string)
+	// waitFuncs are function that needs to be executed after navigation
+	// typically used for waitEvent
+	waitFuncs := make([]func() error, 0)
+
+	// avoid any future panics caused due to go-rod library
+	defer func() {
+		if r := recover(); r != nil {
+			err = errorutil.New("panic on headless action: %v", r)
+		}
+	}()
+
 	for _, act := range actions {
 		switch act.ActionType.ActionType {
 		case ActionNavigate:
 			err = p.NavigateURL(act, outData, variables)
+			if err == nil {
+				// if navigation successful trigger all waitFuncs (if any)
+				for _, waitFunc := range waitFuncs {
+					if waitFunc != nil {
+						if err := waitFunc(); err != nil {
+							return nil, errorutil.NewWithErr(err).Msgf("error occurred while executing waitFunc")
+						}
+					}
+				}
+			}
 		case ActionScript:
 			err = p.RunScript(act, outData)
 		case ActionClick:
@@ -69,7 +91,11 @@ func (p *Page) ExecuteActions(input *contextargs.Context, actions []*Action, var
 		case ActionExtract:
 			err = p.ExtractElement(act, outData)
 		case ActionWaitEvent:
-			err = p.WaitEvent(act, outData)
+			var waitFunc func() error
+			waitFunc, err = p.WaitEvent(act, outData)
+			if waitFunc != nil {
+				waitFuncs = append(waitFuncs, waitFunc)
+			}
 		case ActionFilesInput:
 			if p.options.Options.AllowLocalFileAccess {
 				err = p.FilesInput(act, outData)
@@ -541,38 +567,43 @@ func (p *Page) ExtractElement(act *Action, out map[string]string) error {
 	return nil
 }
 
-type protoEvent struct {
-	event string
-}
-
-// ProtoEvent returns the cdp.Event.Method
-func (p *protoEvent) ProtoEvent() string {
-	return p.event
-}
-
 // WaitEvent waits for an event to happen on the page.
-func (p *Page) WaitEvent(act *Action, out map[string]string /*TODO review unused parameter*/) error {
+func (p *Page) WaitEvent(act *Action, out map[string]string /*TODO review unused parameter*/) (func() error, error) {
 	event := p.getActionArgWithDefaultValues(act, "event")
 	if event == "" {
-		return errors.New("event not recognized")
+		return nil, errors.New("event not recognized")
 	}
-	protoEvent := &protoEvent{event: event}
 
-	// Uses another instance in order to be able to chain the timeout only to the wait operation
-	pageCopy := p.page
-	timeout := p.getActionArg(act, "timeout")
-	if timeout != "" {
-		ts, err := strconv.Atoi(timeout)
+	var waitEvent proto.Event
+	gotType := proto.GetType(event)
+	if gotType == nil {
+		return nil, errorutil.New("event %v does not exist", event)
+	}
+	tmp, ok := reflect.New(gotType).Interface().(proto.Event)
+	if !ok {
+		return nil, errorutil.New("event %v is not a page event", event)
+	}
+	waitEvent = tmp
+	maxDuration := 10 * time.Second // 10 sec is max wait duration for any event
+
+	// allow user to specify max-duration for wait-event
+	if value := p.getActionArgWithDefaultValues(act, "max-duration"); value != "" {
+		var err error
+		maxDuration, err = time.ParseDuration(value)
 		if err != nil {
-			return errors.Wrap(err, "could not get timeout")
-		}
-		if ts > 0 {
-			pageCopy = p.page.Timeout(time.Duration(ts) * time.Second)
+			return nil, errorutil.NewWithErr(err).Msgf("could not parse max-duration")
 		}
 	}
+
 	// Just wait the event to happen
-	pageCopy.WaitEvent(protoEvent)()
-	return nil
+	waitFunc := func() (err error) {
+		// execute actual wait event
+		ctx, cancel := context.WithTimeout(context.Background(), maxDuration)
+		defer cancel()
+		err = contextutil.ExecFunc(ctx, p.page.WaitEvent(waitEvent))
+		return
+	}
+	return waitFunc, nil
 }
 
 // pageElementBy returns a page element from a variety of inputs.
@@ -643,10 +674,6 @@ func selectorBy(selector string) rod.SelectorType {
 	}
 }
 
-func (p *Page) getActionArg(action *Action, arg string) string {
-	return p.getActionArgWithValues(action, arg, nil)
-}
-
 func (p *Page) getActionArgWithDefaultValues(action *Action, arg string) string {
 	return p.getActionArgWithValues(action, arg, generators.MergeMaps(
 		generators.BuildPayloadFromOptions(p.instance.browser.options),