Added integration tests for websocket + misc fixes

integration_tests/websocket/basic.yaml

@@ -0,0 +1,16 @@
+id: basic-request
+
+info:
+  name: Basic Request
+  author: pdteam
+  severity: info
+
+websocket:
+  - address: '{{Scheme}}://{{Hostname}}'
+    inputs:
+      - data: hello
+    matchers:
+      - type: word
+        words:
+          - world
+        part: response

integration_tests/websocket/cswsh.yaml

@@ -0,0 +1,16 @@
+id: basic-cswsh-request
+
+info:
+  name: Basic cswsh Request
+  author: pdteam
+  severity: info
+
+websocket:
+  - address: '{{Scheme}}://{{Hostname}}'
+    headers: 
+      Origin: 'http://evil.com'
+    matchers:
+      - type: word
+        words:
+          - true
+        part: success

integration_tests/websocket/no-cswsh.yaml

@@ -0,0 +1,16 @@
+id: basic-nocswsh-request
+
+info:
+  name: Basic Non-Vulnerable cswsh Request
+  author: pdteam
+  severity: info
+
+websocket:
+  - address: '{{Scheme}}://{{Hostname}}'
+    headers: 
+      Origin: 'http://evil.com'
+    matchers:
+      - type: word
+        words:
+          - true
+        part: success

v2/cmd/integration-test/integration-test.go

@@ -22,10 +22,11 @@ func main() {
 	failed := aurora.Red("[✘]").String()
 
 	protocolTests := map[string]map[string]testutils.TestCase{
-		"http":     httpTestcases,
-		"network":  networkTestcases,
-		"dns":      dnsTestCases,
-		"workflow": workflowTestcases,
+		"http":      httpTestcases,
+		"network":   networkTestcases,
+		"dns":       dnsTestCases,
+		"workflow":  workflowTestcases,
+		"websocket": websocketTestCases,
 	}
 	for proto, tests := range protocolTests {
 		if protocol == "" || protocol == proto {

v2/cmd/integration-test/websocket.go

@@ -0,0 +1,90 @@
+package main
+
+import (
+	"net"
+	"strings"
+
+	"github.com/gobwas/ws/wsutil"
+	"github.com/projectdiscovery/nuclei/v2/internal/testutils"
+)
+
+var websocketTestCases = map[string]testutils.TestCase{
+	"websocket/basic.yaml":    &websocketBasic{},
+	"websocket/cswsh.yaml":    &websocketCswsh{},
+	"websocket/no-cswsh.yaml": &websocketNoCswsh{},
+}
+
+type websocketBasic struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *websocketBasic) Execute(filePath string) error {
+	connHandler := func(conn net.Conn) {
+		for {
+			msg, op, _ := wsutil.ReadClientData(conn)
+			if string(msg) != string("hello") {
+				return
+			}
+			_ = wsutil.WriteServerMessage(conn, op, []byte("world"))
+		}
+	}
+	originValidate := func(origin string) bool {
+		return true
+	}
+	ts := testutils.NewWebsocketServer(connHandler, originValidate)
+	defer ts.Close()
+
+	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
+	if err != nil {
+		return err
+	}
+	if len(results) != 1 {
+		return errIncorrectResultsCount(results)
+	}
+	return nil
+}
+
+type websocketCswsh struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *websocketCswsh) Execute(filePath string) error {
+	connHandler := func(conn net.Conn) {
+
+	}
+	originValidate := func(origin string) bool {
+		return true
+	}
+	ts := testutils.NewWebsocketServer(connHandler, originValidate)
+	defer ts.Close()
+
+	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
+	if err != nil {
+		return err
+	}
+	if len(results) != 1 {
+		return errIncorrectResultsCount(results)
+	}
+	return nil
+}
+
+type websocketNoCswsh struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *websocketNoCswsh) Execute(filePath string) error {
+	connHandler := func(conn net.Conn) {
+
+	}
+	originValidate := func(origin string) bool {
+		return origin == "https://google.com"
+	}
+	ts := testutils.NewWebsocketServer(connHandler, originValidate)
+	defer ts.Close()
+
+	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
+	if err != nil {
+		return err
+	}
+	if len(results) != 0 {
+		return errIncorrectResultsCount(results)
+	}
+	return nil
+}

v2/go.mod

@@ -100,6 +100,7 @@ require (
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/projectdiscovery/blackrock v0.0.0-20210415162320-b38689ae3a2e // indirect
 	github.com/projectdiscovery/cryptoutil v0.0.0-20210805184155-b5d2512f9345 // indirect

v2/go.sum

@@ -556,6 +556,8 @@ github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FI
 github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/performancecopilot/speed v3.0.0+incompatible/go.mod h1:/CLtqpZ5gBg1M9iaPbIdPPGyKcA8hKdoy6hAWba7Yac=
+github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2 h1:JhzVVoYvbOACxoUmOs6V/G4D5nPVUW73rKvXxP4XUJc=
+github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2/go.mod h1:iIss55rKnNBTvrwdmkUpLnDpZoAHvWaiq5+iMmen4AE=
 github.com/pierrec/lz4 v1.0.2-0.20190131084431-473cd7ce01a1/go.mod h1:3/3N9NVKO0jef7pBehbT1qWhCMrIgbYNnFAZCqQ5LRc=
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
 github.com/pingcap/errors v0.11.4/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTwv5KuLY8=

v2/internal/testutils/integration.go

@@ -4,10 +4,14 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"os/exec"
 	"regexp"
 	"strings"
+
+	"github.com/gobwas/ws"
 )
 
 // RunNucleiTemplateAndGetResults returns a list of results for a template
@@ -113,3 +117,23 @@ func NewTCPServer(handler func(conn net.Conn), port ...int) *TCPServer {
 func (s *TCPServer) Close() {
 	s.listener.Close()
 }
+
+// NewWebsocketServer creates a new websocket server from a handler
+func NewWebsocketServer(handler func(conn net.Conn), originValidate func(origin string) bool, port ...int) *httptest.Server {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if value := r.Header.Get("Origin"); value != "" && !originValidate(value) {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+		conn, _, _, err := ws.UpgradeHTTP(r, w)
+		if err != nil {
+			return
+		}
+		go func() {
+			defer conn.Close()
+
+			handler(conn)
+		}()
+	}))
+	return ts
+}

v2/pkg/protocols/others/websocket/websocket.go

@@ -99,7 +99,7 @@ func (r *Request) Compile(options *protocols.ExecuterOptions) error {
 	if len(r.Payloads) > 0 {
 		attackType := r.AttackType
 		if attackType == "" {
-			attackType = "sniper"
+			attackType = "batteringram"
 		}
 		r.attackType = generators.StringToType[attackType]
 
@@ -183,10 +183,10 @@ func (r *Request) executeRequestWithPayloads(input, hostname string, dynamicValu
 	if err != nil {
 		return errors.Wrap(err, "could not parse input url")
 	}
-	payloadValues["Address"] = parsed.Host
+	payloadValues["Hostname"] = parsed.Host
+	payloadValues["Host"] = parsed.Hostname()
 	payloadValues["Scheme"] = parsed.Scheme
 	payloadValues["Path"] = parsed.Path
-	payloadValues["hostname"] = parsed.Hostname()
 
 	for key, value := range r.Headers {
 		finalData, dataErr := expressions.EvaluateByte([]byte(value), payloadValues)
@@ -210,10 +210,14 @@ func (r *Request) executeRequestWithPayloads(input, hostname string, dynamicValu
 		r.options.Progress.IncrementFailedRequestsBy(1)
 		return errors.Wrap(dataErr, "could not evaluate template expressions")
 	}
+	addressToDial := string(finalAddress)
+	if parsed.Path != "" && parsed.Path != "/" {
+		addressToDial = addressToDial + parsed.Path
+	}
 
-	conn, readBuffer, _, err := websocketDialer.Dial(context.Background(), string(finalAddress))
+	conn, readBuffer, _, err := websocketDialer.Dial(context.Background(), addressToDial)
 	if err != nil {
-		r.options.Output.Request(r.options.TemplateID, input, "ssl", err)
+		r.options.Output.Request(r.options.TemplateID, input, "websocket", err)
 		r.options.Progress.IncrementFailedRequestsBy(1)
 		return errors.Wrap(err, "could not connect to server")
 	}
@@ -296,7 +300,7 @@ func (r *Request) executeRequestWithPayloads(input, hostname string, dynamicValu
 	data["ip"] = r.dialer.GetDialedIP(hostname)
 
 	event := eventcreator.CreateEventWithAdditionalOptions(r, data, r.options.Options.Debug || r.options.Options.DebugResponse, func(internalWrappedEvent *output.InternalWrappedEvent) {
-		internalWrappedEvent.OperatorsResult.PayloadValues = payloadValues
+		internalWrappedEvent.OperatorsResult.PayloadValues = dynamicValues
 	})
 	if r.options.Options.Debug || r.options.Options.DebugResponse {
 		responseOutput := responseBuilder.String()

v2/pkg/templates/templates.go

@@ -90,6 +90,8 @@ var TemplateTypes = []string{
 	"headless",
 	"network",
 	"workflow",
+	"ssl",
+	"websocket",
 }
 
 // Type returns the type of the template
@@ -107,6 +109,10 @@ func (t *Template) Type() string {
 		return "network"
 	case len(t.Workflows) > 0:
 		return "workflow"
+	case len(t.RequestsSSL) > 0:
+		return "ssl"
+	case len(t.RequestsWebsocket) > 0:
+		return "websocket"
 	default:
 		return ""
 	}