1 files changed, 244 insertions, 0 deletions

diff --git a/relay/websocket/websocket.go b/relay/websocket/websocket.go
new file mode 100644
index 0000000..cfc3289
--- /dev/null
+++ b/relay/websocket/websocket.go
@@ -0,0 +1,244 @@
+// Package websocket implements RFC 6455 WebSocket framing without external dependencies.
+// Adapted from muxstr's websocket implementation.
+package websocket
+
+import (
+        "bufio"
+        "context"
+        "crypto/rand"
+        "crypto/sha1"
+        "encoding/base64"
+        "encoding/binary"
+        "fmt"
+        "io"
+        "net"
+        "net/http"
+        "strings"
+        "sync"
+        "time"
+)
+
+const (
+        opContinuation = 0x0
+        opBinary       = 0x2
+        opClose        = 0x8
+        opPing         = 0x9
+        opPong         = 0xA
+)
+
+// Conn is a WebSocket connection.
+type Conn struct {
+        rwc    net.Conn
+        br     *bufio.Reader
+        client bool
+        mu     sync.Mutex
+}
+
+func mask(key [4]byte, data []byte) {
+        for i := range data {
+                data[i] ^= key[i%4]
+        }
+}
+
+func (c *Conn) writeFrame(opcode byte, payload []byte) error {
+        c.mu.Lock()
+        defer c.mu.Unlock()
+
+        length := len(payload)
+        header := []byte{0x80 | opcode, 0} // FIN + opcode
+
+        if c.client {
+                header[1] = 0x80 // mask bit
+        }
+
+        switch {
+        case length <= 125:
+                header[1] |= byte(length)
+        case length <= 65535:
+                header[1] |= 126
+                ext := make([]byte, 2)
+                binary.BigEndian.PutUint16(ext, uint16(length))
+                header = append(header, ext...)
+        default:
+                header[1] |= 127
+                ext := make([]byte, 8)
+                binary.BigEndian.PutUint64(ext, uint64(length))
+                header = append(header, ext...)
+        }
+
+        if c.client {
+                var key [4]byte
+                rand.Read(key[:])
+                header = append(header, key[:]...)
+                // mask a copy so we don't modify the caller's slice
+                masked := make([]byte, len(payload))
+                copy(masked, payload)
+                mask(key, masked)
+                payload = masked
+        }
+
+        if _, err := c.rwc.Write(header); err != nil {
+                return err
+        }
+        _, err := c.rwc.Write(payload)
+        return err
+}
+
+func (c *Conn) readFrame() (fin bool, opcode byte, payload []byte, err error) {
+        var hdr [2]byte
+        if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
+                return
+        }
+
+        fin = hdr[0]&0x80 != 0
+        opcode = hdr[0] & 0x0F
+        masked := hdr[1]&0x80 != 0
+        length := uint64(hdr[1] & 0x7F)
+
+        switch length {
+        case 126:
+                var ext [2]byte
+                if _, err = io.ReadFull(c.br, ext[:]); err != nil {
+                        return
+                }
+                length = uint64(binary.BigEndian.Uint16(ext[:]))
+        case 127:
+                var ext [8]byte
+                if _, err = io.ReadFull(c.br, ext[:]); err != nil {
+                        return
+                }
+                length = binary.BigEndian.Uint64(ext[:])
+        }
+
+        var key [4]byte
+        if masked {
+                if _, err = io.ReadFull(c.br, key[:]); err != nil {
+                        return
+                }
+        }
+
+        payload = make([]byte, length)
+        if _, err = io.ReadFull(c.br, payload); err != nil {
+                return
+        }
+
+        if masked {
+                mask(key, payload)
+        }
+        return
+}
+
+// Read reads the next complete message from the connection.
+// It handles ping frames automatically by sending pong responses.
+// It respects context cancellation by setting a read deadline.
+func (c *Conn) Read(ctx context.Context) ([]byte, error) {
+        stop := context.AfterFunc(ctx, func() {
+                c.rwc.SetReadDeadline(time.Now())
+        })
+        defer stop()
+
+        var buf []byte
+        for {
+                fin, opcode, payload, err := c.readFrame()
+                if err != nil {
+                        if ctx.Err() != nil {
+                                return nil, ctx.Err()
+                        }
+                        return nil, err
+                }
+
+                switch opcode {
+                case opPing:
+                        c.writeFrame(opPong, payload)
+                        continue
+                case opClose:
+                        return nil, fmt.Errorf("websocket: close frame received")
+                case opBinary, opContinuation:
+                        buf = append(buf, payload...)
+                        if fin {
+                                return buf, nil
+                        }
+                default:
+                        // text or other opcodes — treat payload as binary
+                        buf = append(buf, payload...)
+                        if fin {
+                                return buf, nil
+                        }
+                }
+        }
+}
+
+// Write sends a binary frame to the connection.
+func (c *Conn) Write(data []byte) error {
+        return c.writeFrame(opBinary, data)
+}
+
+// Ping sends a WebSocket ping frame.
+func (c *Conn) Ping() error {
+        return c.writeFrame(opPing, nil)
+}
+
+// Close sends a close frame with the given code and reason, then closes the
+// underlying connection.
+func (c *Conn) Close(code uint16, reason string) error {
+        payload := make([]byte, 2+len(reason))
+        binary.BigEndian.PutUint16(payload, code)
+        copy(payload[2:], reason)
+        c.writeFrame(opClose, payload)
+        return c.rwc.Close()
+}
+
+// CloseConn closes the underlying network connection without sending a close frame.
+func (c *Conn) CloseConn() error {
+        return c.rwc.Close()
+}
+
+// SetReadDeadline sets the read deadline on the underlying connection.
+func (c *Conn) SetReadDeadline(t time.Time) error {
+        return c.rwc.SetReadDeadline(t)
+}
+
+var wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+
+func acceptKey(key string) string {
+        h := sha1.New()
+        h.Write([]byte(key))
+        h.Write([]byte(wsGUID))
+        return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+// Accept performs the server-side WebSocket handshake, hijacking the HTTP
+// connection and returning a Conn ready for framed I/O.
+func Accept(w http.ResponseWriter, r *http.Request) (*Conn, error) {
+        if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
+                return nil, fmt.Errorf("websocket: missing Upgrade header")
+        }
+
+        key := r.Header.Get("Sec-WebSocket-Key")
+        if key == "" {
+                return nil, fmt.Errorf("websocket: missing Sec-WebSocket-Key")
+        }
+
+        hj, ok := w.(http.Hijacker)
+        if !ok {
+                return nil, fmt.Errorf("websocket: response does not support hijacking")
+        }
+
+        rwc, brw, err := hj.Hijack()
+        if err != nil {
+                return nil, err
+        }
+
+        accept := acceptKey(key)
+        respStr := "HTTP/1.1 101 Switching Protocols\r\n" +
+                "Upgrade: websocket\r\n" +
+                "Connection: Upgrade\r\n" +
+                "Sec-WebSocket-Accept: " + accept + "\r\n\r\n"
+
+        if _, err := rwc.Write([]byte(respStr)); err != nil {
+                rwc.Close()
+                return nil, err
+        }
+
+        return &Conn{rwc: rwc, br: brw.Reader, client: false}, nil
+}