web-vnc/internal/relay/websocket.go

// Package relay bridges a noVNC WebSocket client to a raw TCP VNC server,
// transparently carrying the RFB byte-stream in both directions.
//
// This file implements a minimal RFC 6455 WebSocket server using only the
// Go standard library, tailored to the needs of noVNC: binary message
// frames in both directions, with ping/pong and close handling.
package relay

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// wsConn wraps a hijacked HTTP connection as a WebSocket (server side).
type wsConn struct {
	nc net.Conn
	br *bufio.Reader
	bw *bufio.Writer
}

// handshake performs the WebSocket upgrade and returns a wsConn.
func handshake(w http.ResponseWriter, r *http.Request) (*wsConn, error) {
	if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
		!headerContains(r.Header, "Connection", "upgrade") {
		return nil, errors.New("not a websocket upgrade request")
	}
	key := r.Header.Get("Sec-WebSocket-Key")
	if key == "" {
		return nil, errors.New("missing Sec-WebSocket-Key")
	}

	hj, ok := w.(http.Hijacker)
	if !ok {
		return nil, errors.New("response writer does not support hijacking")
	}
	nc, brw, err := hj.Hijack()
	if err != nil {
		return nil, err
	}

	accept := wsAcceptKey(key)
	_, _ = fmt.Fprintf(brw, "HTTP/1.1 101 Switching Protocols\r\n")
	_, _ = fmt.Fprintf(brw, "Upgrade: websocket\r\n")
	_, _ = fmt.Fprintf(brw, "Connection: Upgrade\r\n")
	_, _ = fmt.Fprintf(brw, "Sec-WebSocket-Accept: %s\r\n", accept)
	_, _ = fmt.Fprintf(brw, "\r\n")
	if err := brw.Flush(); err != nil {
		_ = nc.Close()
		return nil, err
	}
	return &wsConn{nc: nc, br: brw.Reader, bw: brw.Writer}, nil
}

func wsAcceptKey(key string) string {
	h := sha1.New()
	h.Write([]byte(key + wsGUID))
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func headerContains(h http.Header, name, value string) bool {
	for _, v := range h[http.CanonicalHeaderKey(name)] {
		for _, tok := range strings.Split(v, ",") {
			if strings.EqualFold(strings.TrimSpace(tok), value) {
				return true
			}
		}
	}
	return false
}

// wsFrame opcodes
const (
	opContinuation = 0x0
	opText         = 0x1
	opBinary       = 0x2
	opClose        = 0x8
	opPing         = 0x9
	opPong         = 0xA
)

// readFrame reads one WebSocket frame from the client.
// It returns the opcode and (de-masked) payload. Control frames keep their
// own opcode; data frames (text/binary/continuation) are returned as-is and
// the caller is responsible for stream semantics.
func (c *wsConn) readFrame() (opcode byte, payload []byte, err error) {
	var hdr [2]byte
	if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
		return 0, nil, err
	}
	opcode = hdr[0] & 0x0F
	masked := hdr[1]&0x80 != 0
	length := int64(hdr[1] & 0x7F)

	switch length {
	case 126:
		var ext [2]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return 0, nil, err
		}
		length = int64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return 0, nil, err
		}
		length = int64(binary.BigEndian.Uint64(ext[:]))
	}

	var mask [4]byte
	if masked {
		if _, err = io.ReadFull(c.br, mask[:]); err != nil {
			return 0, nil, err
		}
	}

	if length < 0 {
		return 0, nil, errors.New("invalid frame length")
	}
	payload = make([]byte, length)
	if length > 0 {
		if _, err = io.ReadFull(c.br, payload); err != nil {
			return 0, nil, err
		}
	}
	if masked {
		for i := range payload {
			payload[i] ^= mask[i%4]
		}
	}
	return opcode, payload, nil
}

// writeBinary sends a single binary frame to the client (server frames are unmasked).
func (c *wsConn) writeBinary(p []byte) error {
	var hdr []byte
	n := len(p)
	switch {
	case n < 126:
		hdr = []byte{0x82, byte(n)}
	case n < 65536:
		hdr = make([]byte, 4)
		hdr[0] = 0x82
		hdr[1] = 126
		binary.BigEndian.PutUint16(hdr[2:], uint16(n))
	default:
		hdr = make([]byte, 10)
		hdr[0] = 0x82
		hdr[1] = 127
		binary.BigEndian.PutUint64(hdr[2:], uint64(n))
	}
	if _, err := c.bw.Write(hdr); err != nil {
		return err
	}
	if _, err := c.bw.Write(p); err != nil {
		return err
	}
	return c.bw.Flush()
}

// writePong sends a pong frame with the given payload.
func (c *wsConn) writePong(p []byte) error {
	if len(p) > 125 {
		return errors.New("control frame payload too large")
	}
	hdr := []byte{0x8A, byte(len(p))}
	if _, err := c.bw.Write(hdr); err != nil {
		return err
	}
	if _, err := c.bw.Write(p); err != nil {
		return err
	}
	return c.bw.Flush()
}

// writeClose sends a close frame and flushes.
func (c *wsConn) writeClose() {
	_, _ = c.bw.Write([]byte{0x88, 0x00})
	_ = c.bw.Flush()
}

// close closes the underlying connection.
func (c *wsConn) close() error { return c.nc.Close() }

// setWriteDeadline forwards to the underlying connection.
func (c *wsConn) setWriteDeadline(t time.Time) error { return c.nc.SetWriteDeadline(t) }