407 lines
11 KiB
Go
407 lines
11 KiB
Go
package qq_obv11
|
|
|
|
// This file implements a minimal RFC 6455 WebSocket (text/binary frames,
|
|
// ping/pong, close, per-message client masking) using only the standard
|
|
// library. The project pins no third-party WebSocket dependency (go.mod must
|
|
// not change), so the OneBot v11 adapter carries its own tiny frame codec. It
|
|
// is deliberately narrow: it supports what OneBot v11 needs — UTF-8 JSON text
|
|
// frames carried over a single upgraded TCP connection — and nothing more.
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha1"
|
|
"crypto/tls"
|
|
"encoding/base64"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// WebSocket control/data opcodes (RFC 6455 section 5.2).
|
|
const (
|
|
opContinuation = 0x0
|
|
opText = 0x1
|
|
opBinary = 0x2
|
|
opClose = 0x8
|
|
opPing = 0x9
|
|
opPong = 0xA
|
|
)
|
|
|
|
// wsGUID is the fixed handshake GUID from RFC 6455 section 1.3.
|
|
const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
|
|
|
|
// maxFramePayload caps a single WebSocket frame so a malicious peer cannot make
|
|
// the adapter buffer an unbounded payload. Chat/JSON payloads are tiny; this is
|
|
// generous.
|
|
const maxFramePayload = 10 << 20 // 10 MiB
|
|
|
|
// DialContextFunc dials a raw TCP (or TLS) connection. It is injectable so
|
|
// tests can route the handshake through a fake server; the default is a
|
|
// net.Dialer.
|
|
type DialContextFunc func(ctx context.Context, network, address string) (net.Conn, error)
|
|
|
|
// wsConn is a single WebSocket connection over an upgraded TCP stream. It is
|
|
// safe for one concurrent reader plus concurrent writers (writes are mutex
|
|
// serialized); exactly one goroutine may call ReadMessage at a time.
|
|
type wsConn struct {
|
|
conn net.Conn
|
|
br *bufio.Reader
|
|
bw *bufio.Writer
|
|
isClient bool
|
|
writeMu sync.Mutex
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
// newWSConn wraps an upgraded connection. br and bw are the buffered reader and
|
|
// writer from the handshake; passing nil causes them to be created from conn.
|
|
// isClient reports whether this side is the WebSocket client (which must mask
|
|
// its outgoing frames).
|
|
func newWSConn(conn net.Conn, br *bufio.Reader, bw *bufio.Writer, isClient bool) *wsConn {
|
|
if br == nil {
|
|
br = bufio.NewReader(conn)
|
|
}
|
|
if bw == nil {
|
|
bw = bufio.NewWriter(conn)
|
|
}
|
|
return &wsConn{conn: conn, br: br, bw: bw, isClient: isClient}
|
|
}
|
|
|
|
// Close closes the underlying connection exactly once.
|
|
func (c *wsConn) Close() error {
|
|
var err error
|
|
c.closeOnce.Do(func() { err = c.conn.Close() })
|
|
return err
|
|
}
|
|
|
|
// WriteMessage writes a single unmasked/known-masked data or control frame.
|
|
// The caller must hold no locks.
|
|
func (c *wsConn) WriteMessage(opcode byte, payload []byte) error {
|
|
c.writeMu.Lock()
|
|
defer c.writeMu.Unlock()
|
|
|
|
b0 := byte(0x80) | opcode // FIN set on the single frame
|
|
if err := c.bw.WriteByte(b0); err != nil {
|
|
return err
|
|
}
|
|
|
|
maskBit := byte(0)
|
|
if c.isClient {
|
|
maskBit = 0x80
|
|
}
|
|
if err := writeFrameHeader(c.bw, maskBit, len(payload)); err != nil {
|
|
return err
|
|
}
|
|
|
|
var maskKey [4]byte
|
|
if c.isClient {
|
|
if _, err := rand.Read(maskKey[:]); err != nil {
|
|
return err
|
|
}
|
|
if _, err := c.bw.Write(maskKey[:]); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if !c.isClient {
|
|
if _, err := c.bw.Write(payload); err != nil {
|
|
return err
|
|
}
|
|
return c.bw.Flush()
|
|
}
|
|
if err := writeMaskedPayload(c.bw, payload, maskKey); err != nil {
|
|
return err
|
|
}
|
|
return c.bw.Flush()
|
|
}
|
|
|
|
func writeFrameHeader(w *bufio.Writer, maskBit byte, n int) error {
|
|
switch {
|
|
case n < 126:
|
|
return w.WriteByte(maskBit | byte(n))
|
|
case n <= 0xFFFF:
|
|
if err := w.WriteByte(maskBit | 126); err != nil {
|
|
return err
|
|
}
|
|
var ext [2]byte
|
|
binary.BigEndian.PutUint16(ext[:], uint16(n))
|
|
_, err := w.Write(ext[:])
|
|
return err
|
|
default:
|
|
if err := w.WriteByte(maskBit | 127); err != nil {
|
|
return err
|
|
}
|
|
var ext [8]byte
|
|
binary.BigEndian.PutUint64(ext[:], uint64(n))
|
|
_, err := w.Write(ext[:])
|
|
return err
|
|
}
|
|
}
|
|
|
|
func writeMaskedPayload(w *bufio.Writer, payload []byte, maskKey [4]byte) error {
|
|
var buf [4096]byte
|
|
for off := 0; off < len(payload); off += len(buf) {
|
|
end := off + len(buf)
|
|
if end > len(payload) {
|
|
end = len(payload)
|
|
}
|
|
for i, b := range payload[off:end] {
|
|
buf[i] = b ^ maskKey[(off+i)&3]
|
|
}
|
|
if _, err := w.Write(buf[:end-off]); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ReadMessage reads the next complete data message, reassembling fragmented
|
|
// frames, replying to pings, and terminating (io.EOF) on a close frame or
|
|
// network error. Only one goroutine may call it.
|
|
func (c *wsConn) ReadMessage() (byte, []byte, error) {
|
|
var (
|
|
payload []byte
|
|
dataOp byte
|
|
started bool
|
|
)
|
|
|
|
for {
|
|
fin, opcode, data, err := c.readFrame()
|
|
if err != nil {
|
|
return 0, nil, err
|
|
}
|
|
|
|
switch opcode {
|
|
case opPing:
|
|
if err := c.WriteMessage(opPong, data); err != nil {
|
|
return 0, nil, err
|
|
}
|
|
continue
|
|
case opPong:
|
|
continue
|
|
case opClose:
|
|
_ = c.WriteMessage(opClose, nil)
|
|
return opClose, nil, io.EOF
|
|
case opContinuation:
|
|
if !started {
|
|
return 0, nil, errors.New("websocket: unexpected continuation frame")
|
|
}
|
|
payload = append(payload, data...)
|
|
if fin {
|
|
return dataOp, payload, nil
|
|
}
|
|
case opText, opBinary:
|
|
if started {
|
|
return 0, nil, errors.New("websocket: new data frame during fragmented message")
|
|
}
|
|
dataOp = opcode
|
|
if fin {
|
|
return opcode, data, nil
|
|
}
|
|
started = true
|
|
payload = append([]byte(nil), data...)
|
|
default:
|
|
return 0, nil, fmt.Errorf("websocket: unsupported opcode %d", opcode)
|
|
}
|
|
}
|
|
}
|
|
|
|
// readFrame reads a single frame header and payload, unmasking when required.
|
|
func (c *wsConn) 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[:])
|
|
}
|
|
if length > maxFramePayload {
|
|
err = fmt.Errorf("websocket: frame too large (%d bytes)", length)
|
|
return
|
|
}
|
|
|
|
var maskKey [4]byte
|
|
if masked {
|
|
if _, err = io.ReadFull(c.br, maskKey[:]); err != nil {
|
|
return
|
|
}
|
|
}
|
|
|
|
payload = make([]byte, length)
|
|
if _, err = io.ReadFull(c.br, payload); err != nil {
|
|
return
|
|
}
|
|
if masked {
|
|
for i := range payload {
|
|
payload[i] ^= maskKey[i&3]
|
|
}
|
|
}
|
|
return fin, opcode, payload, nil
|
|
}
|
|
|
|
// computeAccept derives the Sec-WebSocket-Accept value for a handshake key.
|
|
func computeAccept(key string) string {
|
|
h := sha1.New()
|
|
_, _ = io.WriteString(h, key+wsGUID)
|
|
return base64.StdEncoding.EncodeToString(h.Sum(nil))
|
|
}
|
|
|
|
// randomWSKey returns a random WebSocket handshake key.
|
|
func randomWSKey() string {
|
|
b := make([]byte, 16)
|
|
_, _ = rand.Read(b)
|
|
return base64.StdEncoding.EncodeToString(b)
|
|
}
|
|
|
|
// upgradeWS performs the server side of the WebSocket handshake by hijacking
|
|
// the HTTP connection and returning the upgraded wsConn.
|
|
func upgradeWS(w http.ResponseWriter, r *http.Request) (*wsConn, error) {
|
|
if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
|
|
!headerHasToken(r.Header, "Connection", "upgrade") {
|
|
http.Error(w, "not a websocket upgrade", http.StatusBadRequest)
|
|
return nil, errors.New("websocket: not an upgrade request")
|
|
}
|
|
key := r.Header.Get("Sec-WebSocket-Key")
|
|
if key == "" {
|
|
http.Error(w, "missing Sec-WebSocket-Key", http.StatusBadRequest)
|
|
return nil, errors.New("websocket: missing handshake key")
|
|
}
|
|
|
|
hj, ok := w.(http.Hijacker)
|
|
if !ok {
|
|
http.Error(w, "hijacking not supported", http.StatusInternalServerError)
|
|
return nil, errors.New("websocket: hijacking not supported")
|
|
}
|
|
netConn, brw, err := hj.Hijack()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("websocket: hijack: %w", err)
|
|
}
|
|
|
|
resp := "HTTP/1.1 101 Switching Protocols\r\n" +
|
|
"Upgrade: websocket\r\n" +
|
|
"Connection: Upgrade\r\n" +
|
|
"Sec-WebSocket-Accept: " + computeAccept(key) + "\r\n\r\n"
|
|
if _, err := brw.WriteString(resp); err != nil {
|
|
_ = netConn.Close()
|
|
return nil, err
|
|
}
|
|
if err := brw.Flush(); err != nil {
|
|
_ = netConn.Close()
|
|
return nil, err
|
|
}
|
|
return newWSConn(netConn, brw.Reader, brw.Writer, false), nil
|
|
}
|
|
|
|
// headerHasToken reports whether a comma-separated header value contains the
|
|
// given token case-insensitively.
|
|
func headerHasToken(h http.Header, name, token string) bool {
|
|
for _, val := range h.Values(name) {
|
|
for _, part := range strings.Split(val, ",") {
|
|
if strings.EqualFold(strings.TrimSpace(part), token) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// dialWS performs the client side of the WebSocket handshake, returning the
|
|
// upgraded wsConn. accessToken, when non-empty, is sent both as an
|
|
// Authorization: Bearer header and as an access_token query parameter for
|
|
// OneBot implementations that expect either.
|
|
func dialWS(ctx context.Context, dial DialContextFunc, rawURL, accessToken string) (*wsConn, error) {
|
|
if dial == nil {
|
|
dial = (&net.Dialer{Timeout: 15 * time.Second}).DialContext
|
|
}
|
|
u, err := url.Parse(rawURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("websocket: parse url %q: %w", rawURL, err)
|
|
}
|
|
if u.Scheme != "ws" && u.Scheme != "wss" {
|
|
return nil, fmt.Errorf("websocket: unsupported scheme %q (want ws/wss)", u.Scheme)
|
|
}
|
|
|
|
host := u.Host
|
|
if u.Port() == "" {
|
|
if u.Scheme == "wss" {
|
|
host += ":443"
|
|
} else {
|
|
host += ":80"
|
|
}
|
|
}
|
|
conn, err := dial(ctx, "tcp", host)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("websocket: dial %s: %w", host, err)
|
|
}
|
|
|
|
if u.Scheme == "wss" {
|
|
tlsConn := tls.Client(conn, &tls.Config{ServerName: u.Hostname()})
|
|
if err := tlsConn.HandshakeContext(ctx); err != nil {
|
|
_ = conn.Close()
|
|
return nil, fmt.Errorf("websocket: tls handshake: %w", err)
|
|
}
|
|
conn = tlsConn
|
|
}
|
|
|
|
brw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
|
|
key := randomWSKey()
|
|
var b strings.Builder
|
|
b.WriteString("GET " + u.RequestURI() + " HTTP/1.1\r\n")
|
|
b.WriteString("Host: " + u.Host + "\r\n")
|
|
b.WriteString("Upgrade: websocket\r\n")
|
|
b.WriteString("Connection: Upgrade\r\n")
|
|
b.WriteString("Sec-WebSocket-Key: " + key + "\r\n")
|
|
b.WriteString("Sec-WebSocket-Version: 13\r\n")
|
|
if accessToken != "" {
|
|
b.WriteString("Authorization: Bearer " + accessToken + "\r\n")
|
|
}
|
|
b.WriteString("\r\n")
|
|
if _, err := brw.WriteString(b.String()); err != nil {
|
|
_ = conn.Close()
|
|
return nil, err
|
|
}
|
|
if err := brw.Flush(); err != nil {
|
|
_ = conn.Close()
|
|
return nil, err
|
|
}
|
|
|
|
resp, err := http.ReadResponse(brw.Reader, &http.Request{Method: "GET"})
|
|
if err != nil {
|
|
_ = conn.Close()
|
|
return nil, fmt.Errorf("websocket: read handshake response: %w", err)
|
|
}
|
|
if resp.StatusCode != http.StatusSwitchingProtocols {
|
|
_ = conn.Close()
|
|
return nil, fmt.Errorf("websocket: unexpected handshake status %d", resp.StatusCode)
|
|
}
|
|
if got := resp.Header.Get("Sec-WebSocket-Accept"); got != computeAccept(key) {
|
|
_ = conn.Close()
|
|
return nil, errors.New("websocket: bad Sec-WebSocket-Accept in response")
|
|
}
|
|
return newWSConn(conn, brw.Reader, brw.Writer, true), nil
|
|
}
|