1 files changed, 159 insertions, 0 deletions
diff --git a/internal/proxy/websocket.go b/internal/proxy/websocket.go
new file mode 100644
index 0000000..10a934c
--- /dev/null
+++ b/internal/proxy/websocket.go
@@ -0,0 +1,159 @@
+// WebSocket interception: after a client's Upgrade: websocket request
+// gets a matching 101 Switching Protocols response back from the origin
+// (see forward's WS branch), the connection stops being HTTP request/
+// response and becomes a long-lived, bidirectional, message-framed
+// stream (RFC 6455) instead. mitmux relays every frame byte-for-byte
+// unmodified in both directions - this is capture, not tampering - while
+// decoding each one's payload for display, recorded to the ws_messages
+// table tagged to the upgrade request's own history entry.
+//
+// Deliberately one row per frame, not per logical message: RFC 6455
+// lets a single message span several frames (opcode 0x0 continuation,
+// FIN unset until the last one), which mitmux does not reassemble.
+// Real-world WebSocket traffic - JSON events, chat messages, game state
+// - is overwhelmingly single-frame; reassembly would need buffering an
+// unbounded number of pending fragmented messages per connection for a
+// case that's rare in practice, which isn't a trade worth making here.
+package proxy
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+)
+
+// WebSocket opcodes (RFC 6455 section 5.2).
+const (
+	wsOpContinuation = 0x0
+	wsOpText         = 0x1
+	wsOpBinary       = 0x2
+	wsOpClose        = 0x8
+	wsOpPing         = 0x9
+	wsOpPong         = 0xa
+)
+
+// isWebSocketUpgradeRequest reports whether r is asking to upgrade to a
+// WebSocket connection (Connection: Upgrade plus Upgrade: websocket).
+// forward() uses this to keep those two headers off stripHopByHop's list
+// for this one request - RFC 7230 correctly treats Connection/Upgrade as
+// hop-by-hop for a normal request, but stripping them here would delete
+// the very signal the origin needs to recognize the upgrade at all,
+// turning every WebSocket connection attempt into a silent 426.
+func isWebSocketUpgradeRequest(r *http.Request) bool {
+	return headerHasToken(r.Header, "Connection", "upgrade") &&
+		strings.EqualFold(r.Header.Get("Upgrade"), "websocket")
+}
+
+// isWebSocketUpgradeResponse reports whether resp is a successful
+// WebSocket upgrade (101 Switching Protocols, with Connection: Upgrade
+// and Upgrade: websocket) - checking the response rather than the
+// request it answers, since a 101 only ever comes back from an origin
+// that accepted the upgrade, and that's the one thing forward() actually
+// needs to know before handing the connection off.
+func isWebSocketUpgradeResponse(resp *http.Response) bool {
+	return resp.StatusCode == http.StatusSwitchingProtocols &&
+		headerHasToken(resp.Header, "Connection", "upgrade") &&
+		strings.EqualFold(resp.Header.Get("Upgrade"), "websocket")
+}
+
+func headerHasToken(h http.Header, name, token string) bool {
+	for _, v := range h.Values(name) {
+		for _, part := range strings.Split(v, ",") {
+			if strings.EqualFold(strings.TrimSpace(part), token) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// relayWSFrame reads exactly one RFC 6455 frame from src, writes the
+// same raw bytes to dst unmodified, and returns the frame's opcode and
+// decoded (unmasked) payload for capture. Masking is direction-
+// dependent - client-to-server frames are always masked, server-to-
+// client frames never are - but relayed bytes are whatever was actually
+// read, so this works correctly regardless of which direction it's
+// called for.
+func relayWSFrame(src io.Reader, dst io.Writer) (opcode byte, payload []byte, err error) {
+	hdr := make([]byte, 2)
+	if _, err = io.ReadFull(src, hdr); err != nil {
+		return 0, nil, err
+	}
+	opcode = hdr[0] & 0x0f
+	masked := hdr[1]&0x80 != 0
+	length := uint64(hdr[1] & 0x7f)
+
+	raw := append([]byte(nil), hdr...)
+
+	switch length {
+	case 126:
+		ext := make([]byte, 2)
+		if _, err = io.ReadFull(src, ext); err != nil {
+			return 0, nil, err
+		}
+		raw = append(raw, ext...)
+		length = uint64(binary.BigEndian.Uint16(ext))
+	case 127:
+		ext := make([]byte, 8)
+		if _, err = io.ReadFull(src, ext); err != nil {
+			return 0, nil, err
+		}
+		raw = append(raw, ext...)
+		length = binary.BigEndian.Uint64(ext)
+	}
+
+	var maskKey [4]byte
+	if masked {
+		if _, err = io.ReadFull(src, maskKey[:]); err != nil {
+			return 0, nil, err
+		}
+		raw = append(raw, maskKey[:]...)
+	}
+
+	// Bounded the same way request/response body capture is (see
+	// maxCaptureBytes) - a length field mitmux doesn't control shouldn't
+	// be able to force an unbounded read/allocation. Relaying (not just
+	// capturing) is refused too: a frame this large is already well
+	// outside normal WebSocket usage, and guessing at a partial relay
+	// would corrupt the stream's framing for whichever side reads next.
+	if length > maxCaptureBytes {
+		return 0, nil, fmt.Errorf("websocket frame too large (%d bytes, over the %d limit)", length, maxCaptureBytes)
+	}
+
+	body := make([]byte, length)
+	if _, err = io.ReadFull(src, body); err != nil {
+		return 0, nil, err
+	}
+	raw = append(raw, body...)
+
+	if _, err = dst.Write(raw); err != nil {
+		return 0, nil, err
+	}
+
+	if !masked {
+		return opcode, body, nil
+	}
+	payload = make([]byte, length)
+	for i := range payload {
+		payload[i] = body[i] ^ maskKey[i%4]
+	}
+	return opcode, payload, nil
+}
+
+// pumpWS relays frames from src to dst until one fails to read/write or
+// a close frame (opcode 0x8) passes through, calling capture with each
+// frame's opcode and decoded payload as it goes.
+func pumpWS(src io.Reader, dst io.Writer, capture func(opcode byte, payload []byte)) {
+	for {
+		opcode, payload, err := relayWSFrame(src, dst)
+		if err != nil {
+			return
+		}
+		capture(opcode, payload)
+		if opcode == wsOpClose {
+			return
+		}
+	}
+}