kjol/go/jsbundler/hmr_ws.go

//go:build dev

package jsbundler

// A minimal RFC 6455 WebSocket server — just enough for one-way server→browser
// push of HMR messages. We hand-roll it (rather than add a dependency) because
// the surface we need is tiny: the handshake, unmasked server text frames, and a
// read loop that answers pings and notices close. No per-message compression, no
// fragmentation, no client→server application data. All of jsbundler's
// HMR support is behind `//go:build dev`, so prod builds compile none of it.

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"sync"
)

// wsGUID is the RFC 6455 magic value concatenated with Sec-WebSocket-Key to
// derive the accept token.
const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

type hub struct {
	mu      sync.Mutex
	clients map[*wsClient]struct{}
}

func newHub() *hub { return &hub{clients: map[*wsClient]struct{}{}} }

type wsClient struct {
	conn net.Conn
	brw  *bufio.ReadWriter
	wmu  sync.Mutex // serialize frame writes across broadcast + pong
}

// ServeWS upgrades an HTTP/1.1 request to a WebSocket and registers the client.
func (h *hub) ServeWS(w http.ResponseWriter, r *http.Request) {
	key := r.Header.Get("Sec-WebSocket-Key")
	if key == "" {
		http.Error(w, "expected a WebSocket handshake", http.StatusBadRequest)
		return
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "connection does not support hijacking", http.StatusInternalServerError)
		return
	}
	conn, brw, err := hj.Hijack()
	if err != nil {
		return
	}

	accept := computeAccept(key)
	if _, err := brw.WriteString(
		"HTTP/1.1 101 Switching Protocols\r\n" +
			"Upgrade: websocket\r\n" +
			"Connection: Upgrade\r\n" +
			"Sec-WebSocket-Accept: " + accept + "\r\n\r\n",
	); err != nil {
		conn.Close()
		return
	}
	if err := brw.Flush(); err != nil {
		conn.Close()
		return
	}

	c := &wsClient{conn: conn, brw: brw}
	h.mu.Lock()
	h.clients[c] = struct{}{}
	h.mu.Unlock()

	go c.readLoop(h)
}

// broadcast writes a text frame to every connected client, dropping any that
// error (disconnected tab).
func (h *hub) broadcast(payload []byte) {
	h.mu.Lock()
	clients := make([]*wsClient, 0, len(h.clients))
	for c := range h.clients {
		clients = append(clients, c)
	}
	h.mu.Unlock()

	for _, c := range clients {
		if err := c.writeText(payload); err != nil {
			h.drop(c)
		}
	}
}

// broadcastJSON marshals v and broadcasts it as a text frame.
func (h *hub) broadcastJSON(v any) {
	b, err := json.Marshal(v)
	if err != nil {
		fmt.Fprintf(os.Stderr, "devhmr: marshal ws message: %v\n", err)
		return
	}
	h.broadcast(b)
}

func (h *hub) drop(c *wsClient) {
	h.mu.Lock()
	if _, ok := h.clients[c]; ok {
		delete(h.clients, c)
		c.conn.Close()
	}
	h.mu.Unlock()
}

// clientCount reports how many browsers are currently connected.
func (h *hub) clientCount() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.clients)
}

// readLoop consumes client frames only to answer pings and to notice a close or
// dead connection, at which point the client is dropped. Application data from
// the client is ignored — this channel is server→browser only.
func (c *wsClient) readLoop(h *hub) {
	defer h.drop(c)
	for {
		op, payload, err := readFrame(c.brw.Reader)
		if err != nil {
			return
		}
		switch op {
		case opClose:
			c.writeFrame(opClose, payload)
			return
		case opPing:
			if c.writeFrame(opPong, payload) != nil {
				return
			}
		}
	}
}

func (c *wsClient) writeText(payload []byte) error { return c.writeFrame(opText, payload) }

// opcodes we handle.
const (
	opText  byte = 0x1
	opClose byte = 0x8
	opPing  byte = 0x9
	opPong  byte = 0xA
)

// writeFrame writes a single unmasked, unfragmented frame (server frames must
// not be masked). Writes are serialized so a broadcast and a pong can't interleave.
func (c *wsClient) writeFrame(opcode byte, payload []byte) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()

	header := make([]byte, 0, 10)
	header = append(header, 0x80|opcode) // FIN + opcode
	n := len(payload)
	switch {
	case n <= 125:
		header = append(header, byte(n))
	case n <= 0xFFFF:
		header = append(header, 126, byte(n>>8), byte(n))
	default:
		header = append(header, 127)
		for i := 7; i >= 0; i-- {
			header = append(header, byte(n>>(8*i)))
		}
	}
	if _, err := c.brw.Write(header); err != nil {
		return err
	}
	if _, err := c.brw.Write(payload); err != nil {
		return err
	}
	return c.brw.Flush()
}

// readFrame reads one frame, unmasking the client payload (client→server frames
// are always masked). Returns the opcode and payload.
func readFrame(r *bufio.Reader) (opcode byte, payload []byte, err error) {
	var h [2]byte
	if _, err = io.ReadFull(r, h[:]); err != nil {
		return
	}
	opcode = h[0] & 0x0F
	masked := h[1]&0x80 != 0
	n := int(h[1] & 0x7F)
	switch n {
	case 126:
		var ext [2]byte
		if _, err = io.ReadFull(r, ext[:]); err != nil {
			return
		}
		n = int(ext[0])<<8 | int(ext[1])
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(r, ext[:]); err != nil {
			return
		}
		n = 0
		for _, b := range ext {
			n = n<<8 | int(b)
		}
	}
	var mask [4]byte
	if masked {
		if _, err = io.ReadFull(r, mask[:]); err != nil {
			return
		}
	}
	payload = make([]byte, n)
	if _, err = io.ReadFull(r, payload); err != nil {
		return
	}
	if masked {
		for i := range payload {
			payload[i] ^= mask[i%4]
		}
	}
	return
}

// computeAccept derives the Sec-WebSocket-Accept response header from the key.
func computeAccept(key string) string {
	s := sha1.Sum([]byte(key + wsGUID))
	return base64.StdEncoding.EncodeToString(s[:])
}