Model websocket read and write functions.

ql/src/semmle/go/frameworks/Websocket.qll

@@ -132,3 +132,150 @@ module WebSocketRequestCall {
     override DataFlow::Node getRequestUrl() { result = this.getArgument(0) }
   }
 }
+
+/*
+ * A message written to a WebSocket, considered as a flow sink for reflected XSS.
+ */
+
+class WebsocketReaderAsSource extends UntrustedFlowSource::Range {
+  WebsocketReaderAsSource() {
+    exists(WebSocketReader r | this = r.getAnOutput().getNode(r.getACall()))
+  }
+}
+
+/**
+ * A function or a method which reads a message from a WebSocket connection.
+ *
+ * Extend this class to refine existing API models. If you want to model new APIs,
+ * extend `WebSocketReader::Range` instead.
+ */
+class WebSocketReader extends Function {
+  WebSocketReader::Range self;
+
+  WebSocketReader() { this = self }
+
+  /** Gets an output of this function that is read from a WebSocket connection. */
+  FunctionOutput getAnOutput() { result = self.getAnOutput() }
+}
+
+/** Provides classes for working with messages read from a WebSocket. */
+module WebSocketReader {
+  /**
+   * A function or a method which reads a message from a WebSocket connection
+   *
+   * Extend this class to model new APIs. If you want to refine existing API models,
+   * extend `WebSocketReader` instead.
+   */
+  abstract class Range extends Function {
+    /**Returns the parameter in which the function stores the message read. */
+    abstract FunctionOutput getAnOutput();
+  }
+
+  /**
+   * Models the `Receive` method of the `golang.org/x/net/websocket` package.
+   */
+  private class GolangXNetCodecRecv extends Range, Method {
+    GolangXNetCodecRecv() {
+      // func (cd Codec) Receive(ws *Conn, v interface{}) (err error)
+      this.hasQualifiedName("golang.org/x/net/websocket", "Codec", "Receive")
+    }
+
+    override FunctionOutput getAnOutput() { result.isParameter(1) }
+  }
+
+  /**
+   * Models the `Read` method of the `golang.org/x/net/websocket` package.
+   */
+  private class GolangXNetConnRead extends Range, Method {
+    GolangXNetConnRead() {
+      // func (ws *Conn) Read(msg []byte) (n int, err error)
+      this.hasQualifiedName("golang.org/x/net/websocket", "Conn", "Read")
+    }
+
+    override FunctionOutput getAnOutput() { result.isParameter(0) }
+  }
+
+  /**
+   * Models the `Read` method of the `nhooyr.io/websocket` package.
+   */
+  private class NhooyrWebsocketRead extends Range, Method {
+    NhooyrWebsocketRead() {
+      // func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error)
+      this.hasQualifiedName("nhooyr.io/websocket", "Conn", "Read")
+    }
+
+    override FunctionOutput getAnOutput() { result.isResult(1) }
+  }
+
+  /**
+   * Models the `Reader` method of the `nhooyr.io/websocket` package.
+   */
+  private class NhooyrWebsocketReader extends Range, Method {
+    NhooyrWebsocketReader() {
+      // func (c *Conn) Reader(ctx context.Context) (MessageType, io.Reader, error)
+      this.hasQualifiedName("nhooyr.io/websocket", "Conn", "Reader")
+    }
+
+    override FunctionOutput getAnOutput() { result.isResult(1) }
+  }
+
+  /**
+   * Models the `ReadFrame`function of the `github.com/gobwas/ws` package.
+   */
+  private class GobwasWsReadFrame extends Range {
+    GobwasWsReadFrame() {
+      // func ReadFrame(r io.Reader) (f Frame, err error)
+      this.hasQualifiedName("github.com/gobwas/ws", "ReadFrame")
+    }
+
+    override FunctionOutput getAnOutput() { result.isResult(0) }
+  }
+
+  /**
+   * Models the `ReadHeader`function of the `github.com/gobwas/ws` package.
+   */
+  private class GobwasWsReadHeader extends Range {
+    GobwasWsReadHeader() {
+      // func ReadHeader(r io.Reader) (h Header, err error)
+      this.hasQualifiedName("github.com/gobwas/ws", "ReadHeader")
+    }
+
+    override FunctionOutput getAnOutput() { result.isResult(0) }
+  }
+
+  /**
+   * Models the `ReadJson` function of the `github.com/gorilla/websocket` package.
+   */
+  private class GorillaWebsocketReadJson extends Range {
+    GorillaWebsocketReadJson() {
+      // func ReadJSON(c *Conn, v interface{}) error
+      this.hasQualifiedName("github.com/gorilla/websocket", "ReadJSON")
+    }
+
+    override FunctionOutput getAnOutput() { result.isParameter(1) }
+  }
+
+  /**
+   * Models the `ReadJson` method of the `github.com/gorilla/websocket` package.
+   */
+  private class GorillaWebsocketConnReadJson extends Range, Method {
+    GorillaWebsocketConnReadJson() {
+      // func (c *Conn) ReadJSON(v interface{}) error
+      this.hasQualifiedName("github.com/gorilla/websocket", "Conn", "ReadJSON")
+    }
+
+    override FunctionOutput getAnOutput() { result.isParameter(0) }
+  }
+
+  /**
+   * Models the `ReadMessage` method of the `github.com/gorilla/websocket` package.
+   */
+  private class GorillaWebsocketReadMessage extends Range, Method {
+    GorillaWebsocketReadMessage() {
+      // func (c *Conn) ReadMessage() (messageType int, p []byte, err error)
+      this.hasQualifiedName("github.com/gorilla/websocket", "Conn", "ReadMessage")
+    }
+
+    override FunctionOutput getAnOutput() { result.isResult(1) }
+  }
+}

ql/test/library-tests/semmle/go/frameworks/Websocket/DialFunction.expected

@@ -1,9 +1,13 @@
-| DialFunction.go:25:2:25:43 | call to Dial | DialFunction.go:25:17:25:30 | untrustedInput |
-| DialFunction.go:28:2:28:29 | call to DialConfig | DialFunction.go:27:35:27:48 | untrustedInput |
-| DialFunction.go:30:2:30:49 | call to Dial | DialFunction.go:30:30:30:43 | untrustedInput |
-| DialFunction.go:33:2:33:33 | call to Dial | DialFunction.go:33:14:33:27 | untrustedInput |
-| DialFunction.go:35:2:35:56 | call to DialContext | DialFunction.go:35:37:35:50 | untrustedInput |
-| DialFunction.go:37:2:37:44 | call to Dial | DialFunction.go:37:30:37:43 | untrustedInput |
-| DialFunction.go:40:2:40:45 | call to Dial | DialFunction.go:40:31:40:44 | untrustedInput |
-| DialFunction.go:42:2:42:31 | call to BuildProxy | DialFunction.go:42:17:42:30 | untrustedInput |
-| DialFunction.go:43:2:43:24 | call to New | DialFunction.go:43:10:43:23 | untrustedInput |
+| DialFunction.go:21:2:21:43 | call to Dial | DialFunction.go:21:17:21:30 | untrustedInput |
+| DialFunction.go:24:2:24:29 | call to DialConfig | DialFunction.go:23:35:23:48 | untrustedInput |
+| DialFunction.go:26:2:26:49 | call to Dial | DialFunction.go:26:30:26:43 | untrustedInput |
+| DialFunction.go:29:2:29:33 | call to Dial | DialFunction.go:29:14:29:27 | untrustedInput |
+| DialFunction.go:31:2:31:56 | call to DialContext | DialFunction.go:31:37:31:50 | untrustedInput |
+| DialFunction.go:33:2:33:44 | call to Dial | DialFunction.go:33:30:33:43 | untrustedInput |
+| DialFunction.go:36:2:36:45 | call to Dial | DialFunction.go:36:31:36:44 | untrustedInput |
+| DialFunction.go:38:2:38:31 | call to BuildProxy | DialFunction.go:38:17:38:30 | untrustedInput |
+| DialFunction.go:39:2:39:24 | call to New | DialFunction.go:39:10:39:23 | untrustedInput |
+| WebsocketReadWrite.go:30:12:30:42 | call to Dial | WebsocketReadWrite.go:30:27:30:29 | uri |
+| WebsocketReadWrite.go:40:14:40:50 | call to Dial | WebsocketReadWrite.go:40:42:40:44 | uri |
+| WebsocketReadWrite.go:50:17:50:37 | call to Dial | WebsocketReadWrite.go:50:29:50:31 | uri |
+| WebsocketReadWrite.go:66:20:66:51 | call to Dial | WebsocketReadWrite.go:66:48:66:50 | uri |

ql/test/library-tests/semmle/go/frameworks/Websocket/DialFunction.go

@@ -1,10 +1,6 @@
 package main
 
-//go:generate depstubber -vendor github.com/gobwas/ws Dialer Dial
-//go:generate depstubber -vendor github.com/gorilla/websocket Dialer
 //go:generate depstubber -vendor github.com/sacOO7/gowebsocket "" New,BuildProxy
-//go:generate depstubber -vendor golang.org/x/net/websocket "" Dial,NewConfig,DialConfig
-//go:generate depstubber -vendor nhooyr.io/websocket "" Dial
 
 import (
 	"context"

ql/test/library-tests/semmle/go/frameworks/Websocket/Read.expected

@@ -0,0 +1,8 @@
+| WebsocketReadWrite.go:31:7:31:10 | definition of xnet |
+| WebsocketReadWrite.go:35:3:35:7 | definition of xnet2 |
+| WebsocketReadWrite.go:41:3:41:40 | ... := ...[1] |
+| WebsocketReadWrite.go:44:3:44:48 | ... := ...[1] |
+| WebsocketReadWrite.go:51:7:51:16 | definition of gorillaMsg |
+| WebsocketReadWrite.go:55:3:55:10 | definition of gorilla2 |
+| WebsocketReadWrite.go:61:3:61:38 | ... := ...[1] |
+| WebsocketReadWrite.go:67:3:67:36 | ... := ...[0] |

ql/test/library-tests/semmle/go/frameworks/Websocket/Read.ql

@@ -0,0 +1,6 @@
+import go
+import semmle.go.frameworks.Websocket
+
+from WebSocketReader r, DataFlow::Node nd
+where nd = r.getAnOutput().getNode(r.getACall())
+select nd

ql/test/library-tests/semmle/go/frameworks/Websocket/WebsocketReadWrite.go

@@ -0,0 +1,74 @@
+package main
+
+//go:generate depstubber -vendor github.com/gobwas/ws Dialer ReadFrame,WriteFrame,NewTextFrame,Dial
+//go:generate depstubber -vendor github.com/gorilla/websocket Dialer ReadJSON,WriteJSON,NewPreparedMessage
+//go:generate depstubber -vendor golang.org/x/net/websocket Codec Dial,NewConfig,DialConfig
+//go:generate depstubber -vendor nhooyr.io/websocket "" Dial
+
+import (
+	"context"
+	"io"
+	"net/http"
+
+	gobwas "github.com/gobwas/ws"
+	gorilla "github.com/gorilla/websocket"
+	websocket "golang.org/x/net/websocket"
+	nhooyr "nhooyr.io/websocket"
+)
+
+func marshal(v interface{}) (data []byte, payloadType byte, err error) {
+	return nil, 0, nil
+}
+func unmarshal(data []byte, payloadType byte, v interface{}) (err error) {
+	return nil
+}
+
+func xss(w http.ResponseWriter, r *http.Request) {
+	uri := r.Header.Get("X-Header")
+	origin := "test"
+	{
+		ws, _ := websocket.Dial(uri, "", origin)
+		var xnet = make([]byte, 512)
+		ws.Read(xnet)
+		ws.Write([]byte(xnet))
+		codec := &websocket.Codec{marshal, unmarshal}
+		xnet2 := make([]byte, 512)
+		codec.Receive(ws, xnet2)
+		codec.Send(ws, []byte(xnet2))
+	}
+	{
+		n, _, _ := nhooyr.Dial(context.TODO(), uri, nil)
+		_, nhooyr, _ := n.Read(context.TODO())
+		n.Write(context.TODO(), 0, nhooyr)
+
+		_, nhooyrReader, _ := n.Reader(context.TODO())
+		writer, _ := n.Writer(context.TODO(), 0)
+		io.Copy(writer, nhooyrReader)
+	}
+	{
+		dialer := gorilla.Dialer{}
+		conn, _, _ := dialer.Dial(uri, nil)
+		var gorillaMsg = make([]byte, 512)
+		gorilla.ReadJSON(conn, gorillaMsg)
+		gorilla.WriteJSON(conn, gorillaMsg)
+
+		gorilla2 := make([]byte, 512)
+		conn.ReadJSON(gorilla2)
+		pm, _ := gorilla.NewPreparedMessage(0, gorilla2)
+		conn.WritePreparedMessage(pm)
+		conn.WriteJSON(gorilla2)
+
+		_, gorilla3, _ := conn.ReadMessage()
+		conn.WriteMessage(0, gorilla3)
+
+	}
+	{
+		conn, _, _, _ := gobwas.Dial(context.TODO(), uri)
+		frame, _ := gobwas.ReadFrame(conn)
+
+		gobwas.WriteFrame(conn, frame)
+
+		resp := gobwas.NewTextFrame([]byte(uri))
+		gobwas.WriteFrame(conn, resp)
+	}
+}

ql/test/library-tests/semmle/go/frameworks/Websocket/go.mod

@@ -1,11 +1,11 @@
-module main
+module codeql-go-tests/frameworks/Websocket
 
 go 1.14
 
 require (
 	github.com/gobwas/ws v1.0.3
 	github.com/gorilla/websocket v1.4.2
 	github.com/sacOO7/gowebsocket v0.0.0-20180719182212-1436bb906a4e
-	golang.org/x/net v0.0.0-20200421231249-e086a090c8fd
+	golang.org/x/net v0.0.0-20200505041828-1ed23360d12c
 	nhooyr.io/websocket v1.8.5
 )