basic implementation of websockets connection, hijacking, transmission from scratch over net/http in Go

websockets.go

package main

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

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

type Frame struct {
	Fin     bool
	Opcode  byte
	Payload []byte
}

func readFrame(r io.Reader) (*Frame, error) {
	var hdr [2]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return nil, err
	}

	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(r, ext[:]); err != nil {
			return nil, err
		}
		length = uint64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err := io.ReadFull(r, ext[:]); err != nil {
			return nil, err
		}
		length = binary.BigEndian.Uint64(ext[:])
	}

	var maskKey [4]byte
	if masked {
		if _, err := io.ReadFull(r, maskKey[:]); err != nil {
			return nil, err
		}
	}

	payload := make([]byte, length)
	if _, err := io.ReadFull(r, payload); err != nil {
		return nil, err
	}

	if masked {
		for i := range payload {
			payload[i] ^= maskKey[i%4]
		}
	}

	return &Frame{Fin: fin, Opcode: opcode, Payload: payload}, nil
}

func writeFrame(w io.Writer, opcode byte, payload []byte) error {
	hdr := []byte{0x80 | opcode}

	n := len(payload)
	switch {
	case n <= 125:
		hdr = append(hdr, byte(n))
	case n <= 0xFFFF:
		hdr = append(hdr, 126, 0, 0)
		binary.BigEndian.PutUint16(hdr[2:], uint16(n))
	default:
		hdr = append(hdr, 127, 0, 0, 0, 0, 0, 0, 0, 0)
		binary.BigEndian.PutUint64(hdr[2:], uint64(n))
	}

	if _, err := w.Write(hdr); err != nil {
		return err
	}
	_, err := w.Write(payload)
	return err
}

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

func handshake(w http.ResponseWriter, r *http.Request) (net.Conn, *bufio.ReadWriter, error) {
	key := r.Header.Get("Sec-WebSocket-Key")
	conn, brw, err := w.(http.Hijacker).Hijack()
	if err != nil {
		return nil, nil, err
	}
	resp := "HTTP/1.1 101 Switching Protocols\r\n" +
		"Upgrade: websocket\r\n" +
		"Connection: Upgrade\r\n" +
		"Sec-WebSocket-Accept: " + acceptKey(key) + "\r\n\r\n"
	brw.WriteString(resp)
	brw.Flush()
	return conn, brw, nil
}

func main() {
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, brw, err := handshake(w, r)
		if err != nil {
			return
		}
		defer conn.Close()

		for {
			f, err := readFrame(brw)
			if err != nil {
				return
			}
			fmt.Printf("fin=%v opcode=0x%x payload=%q\n", f.Fin, f.Opcode, f.Payload)

			switch f.Opcode {
			case 0x1, 0x2:
				writeFrame(brw, f.Opcode, f.Payload)
			case 0x9:
				writeFrame(brw, 0xA, f.Payload)
			case 0x8:
				writeFrame(brw, 0x8, f.Payload)
				brw.Flush()
				return
			}
			brw.Flush()
		}
	})
	http.ListenAndServe(":8080", nil)
}