// 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
		}
	}
}