《MMO 服务器开发》6.1 WebSocket 握手、掩码与分片

// snippet: internal/ws/handshake.go（节选：AcceptKey 与 101 响应）

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

// AcceptKey returns the Sec-WebSocket-Accept value for a client key:
// key + GUID, SHA-1, then Base64.
func AcceptKey(clientKey string) string {
	h := sha1.New()
	_, _ = io.WriteString(h, clientKey+wsGUID)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

// ……中间是请求行解析、Upgrade / Connection / Version 校验……

	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"
	if _, err := io.WriteString(w, resp); err != nil {
		return "", fmt.Errorf("ws: write handshake response: %w", err)
	}
	return path, nil

// file: internal/ws/frame.go
// tier: L1
// verify: go test ./internal/ws/ -run 'TestMask|TestFrame|TestReadMessage|TestControl|TestOversized|TestContinuation|TestUnmasked' -v

package ws

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

// Opcodes from RFC 6455 §5.2. 0x3-0x7 and 0xB-0xF stay reserved.
const (
	OpContinuation byte = 0x0
	OpText         byte = 0x1
	OpBinary       byte = 0x2
	OpClose        byte = 0x8
	OpPing         byte = 0x9
	OpPong         byte = 0xA
)

// Frame errors the connection loop maps onto a close code.
var (
	ErrReservedBits        = errors.New("ws: reserved bits set")
	ErrFragmentedCtl       = errors.New("ws: control frame must not be fragmented")
	ErrControlTooLarge     = errors.New("ws: control frame payload exceeds 125 bytes")
	ErrUnexpectedCont      = errors.New("ws: continuation frame without a start")
	ErrInterleavedFrame    = errors.New("ws: data frame interleaved inside a fragmented message")
	ErrPayloadTooLarge     = errors.New("ws: payload exceeds configured maximum")
	ErrBadOpcode           = errors.New("ws: unknown opcode")
	ErrControlInterleave   = errors.New("ws: control frame interleaved inside a fragmented message")
	ErrUnmaskedClientFrame = errors.New("ws: client frame was not masked")
)

func isControl(op byte) bool { return op&0x08 != 0 }

// Frame is one wire-level frame. Payload is already unmasked when ReadFrame
// returns it; Masked records what the sender did on the wire.
type Frame struct {
	Fin     bool
	Opcode  byte
	Masked  bool
	Payload []byte
}

// Mask applies the 4-byte masking key in place. XOR is an involution, so the
// same function masks and unmasks.
func Mask(key [4]byte, b []byte) {
	for i := range b {
		b[i] ^= key[i&3]
	}
}

// ReadFrame reads exactly one frame. maxPayload caps the declared length;
// 0 means no cap. The masking key is applied whenever the MASK bit is set.
func ReadFrame(r io.Reader, maxPayload uint64) (*Frame, error) {
	var head [2]byte
	if _, err := io.ReadFull(r, head[:]); err != nil {
		return nil, err
	}
	f := &Frame{Fin: head[0]&0x80 != 0, Opcode: head[0] & 0x0F}
	if head[0]&0x70 != 0 {
		return nil, ErrReservedBits
	}
	f.Masked = head[1]&0x80 != 0
	length := uint64(head[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[:])
		if length>>63 != 0 {
			return nil, errors.New("ws: 64-bit payload length has the high bit set")
		}
	}
	if maxPayload > 0 && length > maxPayload {
		return nil, fmt.Errorf("%w: %d > %d", ErrPayloadTooLarge, length, maxPayload)
	}
	var key [4]byte
	if f.Masked {
		if _, err := io.ReadFull(r, key[:]); err != nil {
			return nil, err
		}
	}
	f.Payload = make([]byte, length)
	if _, err := io.ReadFull(r, f.Payload); err != nil {
		return nil, err
	}
	if f.Masked {
		Mask(key, f.Payload)
	}
	return f, nil
}

// WriteFrame serialises one frame. When masked is true a fresh random masking
// key is generated, as the RFC requires of clients.
func WriteFrame(w io.Writer, fin bool, opcode byte, payload []byte, masked bool) error {
	head := make([]byte, 0, 14)
	b0 := opcode & 0x0F
	if fin {
		b0 |= 0x80
	}
	head = append(head, b0)

	maskBit := byte(0)
	if masked {
		maskBit = 0x80
	}
	n := len(payload)
	switch {
	case n < 126:
		head = append(head, maskBit|byte(n))
	case n <= 0xFFFF:
		head = append(head, maskBit|126, byte(n>>8), byte(n))
	default:
		head = append(head, maskBit|127)
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(n))
		head = append(head, ext[:]...)
	}

	var key [4]byte
	if masked {
		if _, err := rand.Read(key[:]); err != nil {
			return err
		}
		head = append(head, key[:]...)
	}
	if _, err := w.Write(head); err != nil {
		return err
	}
	if n == 0 {
		return nil
	}
	if !masked {
		_, err := w.Write(payload)
		return err
	}
	buf := make([]byte, n) // copy: masking must not mutate the caller's slice
	copy(buf, payload)
	Mask(key, buf)
	_, err := w.Write(buf)
	return err
}

// ReadMessage reassembles one complete data message. Control frames are never
// buffered: they come back immediately so ping/pong/close can be answered even
// in the middle of a fragmented data message.
func ReadMessage(r io.Reader, maxPayload uint64) (byte, []byte, error) {
	return ReadMessageMasked(r, maxPayload, false)
}

// ReadMessageMasked is ReadMessage with an explicit masking policy. Servers
// pass requireMask=true: RFC 6455 §5.1 makes an unmasked client frame a
// protocol error.
func ReadMessageMasked(r io.Reader, maxPayload uint64, requireMask bool) (opcode byte, payload []byte, err error) {
	var (
		msgOp byte
		buf   []byte
	)
	for {
		f, err := ReadFrame(r, maxPayload)
		if err != nil {
			return 0, nil, err
		}
		if requireMask && !f.Masked {
			return 0, nil, ErrUnmaskedClientFrame
		}
		if isControl(f.Opcode) {
			if !f.Fin {
				return 0, nil, ErrFragmentedCtl
			}
			if len(f.Payload) > 125 {
				return 0, nil, ErrControlTooLarge
			}
			if msgOp != 0 {
				return 0, nil, ErrControlInterleave
			}
			return f.Opcode, f.Payload, nil
		}
		switch f.Opcode {
		case OpContinuation:
			if msgOp == 0 {
				return 0, nil, ErrUnexpectedCont
			}
			buf = append(buf, f.Payload...)
		case OpText, OpBinary:
			if msgOp != 0 {
				return 0, nil, ErrInterleavedFrame
			}
			msgOp = f.Opcode
			buf = append(buf, f.Payload...)
		default:
			return 0, nil, fmt.Errorf("%w: 0x%x", ErrBadOpcode, f.Opcode)
		}
		if maxPayload > 0 && uint64(len(buf)) > maxPayload {
			return 0, nil, fmt.Errorf("%w: %d > %d", ErrPayloadTooLarge, len(buf), maxPayload)
		}
		if f.Fin {
			return msgOp, buf, nil
		}
	}
}