devplacepy/nodes/client/ws.go

package main

import (
	"bufio"
	"context"
	"crypto/rand"
	"crypto/sha1"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
	"unicode/utf8"
)

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

	wsGUID           = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
	wsWriteTimeout   = 30 * time.Second
	wsHandshakeLimit = 15 * time.Second
	maxControlBytes  = 125
)

type CloseError struct {
	Code   int
	Reason string
}

func (e *CloseError) Error() string {
	return fmt.Sprintf("websocket closed code=%d reason=%q", e.Code, e.Reason)
}

type Conn struct {
	rwc        net.Conn
	br         *bufio.Reader
	client     bool
	maxMessage int
	wmu        sync.Mutex
	closed     bool
}

func newConn(rwc net.Conn, br *bufio.Reader, client bool) *Conn {
	return &Conn{rwc: rwc, br: br, client: client, maxMessage: MaxInboundMessage}
}

func Dial(ctx context.Context, rawURL string) (*Conn, error) {
	target, err := url.Parse(rawURL)
	if err != nil {
		return nil, errors.New("websocket url cannot be parsed")
	}
	secure := false
	switch target.Scheme {
	case "wss", "https":
		secure = true
	case "ws", "http":
	default:
		return nil, errors.New("websocket url scheme must be ws, wss, http or https")
	}
	host := target.Hostname()
	port := target.Port()
	if port == "" {
		port = "80"
		if secure {
			port = "443"
		}
	}
	address := net.JoinHostPort(host, port)
	dialer := &net.Dialer{Timeout: wsHandshakeLimit}
	raw, err := dialer.DialContext(ctx, "tcp", address)
	if err != nil {
		return nil, fmt.Errorf("websocket connect to %s failed: %w", address, stripURL(err))
	}
	var rwc net.Conn = raw
	if secure {
		tlsConn := tls.Client(raw, &tls.Config{ServerName: host, MinVersion: tls.VersionTLS12})
		_ = raw.SetDeadline(time.Now().Add(wsHandshakeLimit))
		if err := tlsConn.HandshakeContext(ctx); err != nil {
			raw.Close()
			return nil, fmt.Errorf("websocket tls handshake failed: %w", err)
		}
		rwc = tlsConn
	}
	_ = rwc.SetDeadline(time.Now().Add(wsHandshakeLimit))
	keyBytes := make([]byte, 16)
	if _, err := rand.Read(keyBytes); err != nil {
		rwc.Close()
		return nil, err
	}
	key := base64.StdEncoding.EncodeToString(keyBytes)
	path := target.RequestURI()
	if path == "" {
		path = "/"
	}
	request := fmt.Sprintf(
		"GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\nUser-Agent: devplace-node/%s\r\n\r\n",
		path, target.Host, key, Version,
	)
	if _, err := io.WriteString(rwc, request); err != nil {
		rwc.Close()
		return nil, fmt.Errorf("websocket handshake write failed: %w", err)
	}
	br := bufio.NewReaderSize(rwc, 16*1024)
	response, err := http.ReadResponse(br, &http.Request{Method: "GET"})
	if err != nil {
		rwc.Close()
		return nil, fmt.Errorf("websocket handshake read failed: %w", err)
	}
	if response.StatusCode != http.StatusSwitchingProtocols {
		rwc.Close()
		return nil, fmt.Errorf("websocket handshake refused with http status %d", response.StatusCode)
	}
	if !strings.EqualFold(response.Header.Get("Upgrade"), "websocket") ||
		response.Header.Get("Sec-WebSocket-Accept") != acceptKey(key) {
		rwc.Close()
		return nil, errors.New("websocket handshake answer is invalid")
	}
	_ = rwc.SetDeadline(time.Time{})
	return newConn(rwc, br, true), nil
}

func stripURL(err error) error {
	var urlErr *url.Error
	if errors.As(err, &urlErr) {
		return urlErr.Err
	}
	return err
}

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

func (c *Conn) SetReadDeadline(t time.Time) error { return c.rwc.SetReadDeadline(t) }

func (c *Conn) Close() error {
	c.wmu.Lock()
	c.closed = true
	c.wmu.Unlock()
	return c.rwc.Close()
}

func (c *Conn) writeFrame(fin bool, op byte, payload []byte) error {
	header := make([]byte, 0, 14)
	first := op
	if fin {
		first |= 0x80
	}
	header = append(header, first)
	maskBit := byte(0)
	if c.client {
		maskBit = 0x80
	}
	size := len(payload)
	switch {
	case size <= 125:
		header = append(header, maskBit|byte(size))
	case size <= 0xFFFF:
		header = append(header, maskBit|126, 0, 0)
		binary.BigEndian.PutUint16(header[len(header)-2:], uint16(size))
	default:
		header = append(header, maskBit|127, 0, 0, 0, 0, 0, 0, 0, 0)
		binary.BigEndian.PutUint64(header[len(header)-8:], uint64(size))
	}
	body := payload
	if c.client {
		mask := make([]byte, 4)
		if _, err := rand.Read(mask); err != nil {
			return err
		}
		header = append(header, mask...)
		body = make([]byte, size)
		for i := range payload {
			body[i] = payload[i] ^ mask[i%4]
		}
	}
	if err := c.rwc.SetWriteDeadline(time.Now().Add(wsWriteTimeout)); err != nil {
		return err
	}
	if _, err := c.rwc.Write(header); err != nil {
		return err
	}
	_, err := c.rwc.Write(body)
	return err
}

func (c *Conn) WriteText(text string) error {
	data := []byte(text)
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.closed {
		return net.ErrClosed
	}
	op := byte(opText)
	for first := true; first || len(data) > 0; first = false {
		chunk := data
		if len(chunk) > MaxWSFramePayload {
			chunk = data[:MaxWSFramePayload]
		}
		data = data[len(chunk):]
		if err := c.writeFrame(len(data) == 0, op, chunk); err != nil {
			return err
		}
		op = opContinuation
	}
	return nil
}

func (c *Conn) writeControl(op byte, payload []byte) error {
	if len(payload) > maxControlBytes {
		payload = payload[:maxControlBytes]
	}
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.closed {
		return net.ErrClosed
	}
	return c.writeFrame(true, op, payload)
}

func (c *Conn) WriteClose(code int, reason string) error {
	if len(reason) > 100 {
		reason = reason[:100]
	}
	payload := make([]byte, 2, 2+len(reason))
	binary.BigEndian.PutUint16(payload, uint16(code))
	payload = append(payload, reason...)
	return c.writeControl(opClose, payload)
}

func (c *Conn) readFrame() (fin bool, op byte, payload []byte, err error) {
	head := make([]byte, 2)
	if _, err = io.ReadFull(c.br, head); err != nil {
		return
	}
	fin = head[0]&0x80 != 0
	if head[0]&0x70 != 0 {
		err = &CloseError{Code: CloseProtocolError, Reason: "reserved bits set"}
		return
	}
	op = head[0] & 0x0F
	masked := head[1]&0x80 != 0
	if masked == c.client {
		err = &CloseError{Code: CloseProtocolError, Reason: "bad masking"}
		return
	}
	length := uint64(head[1] & 0x7F)
	switch length {
	case 126:
		ext := make([]byte, 2)
		if _, err = io.ReadFull(c.br, ext); err != nil {
			return
		}
		length = uint64(binary.BigEndian.Uint16(ext))
	case 127:
		ext := make([]byte, 8)
		if _, err = io.ReadFull(c.br, ext); err != nil {
			return
		}
		length = binary.BigEndian.Uint64(ext)
	}
	if op&0x8 != 0 && (length > maxControlBytes || !fin) {
		err = &CloseError{Code: CloseProtocolError, Reason: "bad control frame"}
		return
	}
	if length > uint64(c.maxMessage) || length > uint64(MaxResultMessage) {
		err = &CloseError{Code: CloseTooLarge, Reason: "frame too large"}
		return
	}
	var mask []byte
	if masked {
		mask = make([]byte, 4)
		if _, err = io.ReadFull(c.br, mask); 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] ^= mask[i%4]
		}
	}
	return
}

func (c *Conn) fail(code int, reason string) error {
	_ = c.WriteClose(code, reason)
	return &CloseError{Code: code, Reason: reason}
}

func (c *Conn) ReadMessage() ([]byte, error) {
	var message []byte
	inProgress := false
	for {
		fin, op, payload, err := c.readFrame()
		if err != nil {
			var closeErr *CloseError
			if errors.As(err, &closeErr) {
				return nil, c.fail(closeErr.Code, closeErr.Reason)
			}
			return nil, err
		}
		switch op {
		case opPing:
			if err := c.writeControl(opPong, payload); err != nil {
				return nil, err
			}
		case opPong:
		case opClose:
			code, reason := CloseNormal, ""
			if len(payload) >= 2 {
				code = int(binary.BigEndian.Uint16(payload))
				reason = string(payload[2:])
			}
			_ = c.WriteClose(code, "")
			return nil, &CloseError{Code: code, Reason: reason}
		case opText, opBinary, opContinuation:
			if op == opContinuation && !inProgress {
				return nil, c.fail(CloseProtocolError, "unexpected continuation")
			}
			if op != opContinuation && inProgress {
				return nil, c.fail(CloseProtocolError, "interleaved data frames")
			}
			if op == opBinary {
				return nil, c.fail(CloseUnsupported, "binary frames are not supported")
			}
			inProgress = !fin
			if len(message)+len(payload) > c.maxMessage {
				return nil, c.fail(CloseTooLarge, "message too large")
			}
			message = append(message, payload...)
			if fin {
				if !utf8.Valid(message) {
					return nil, c.fail(CloseInvalidPayload, "invalid utf-8")
				}
				return message, nil
			}
		default:
			return nil, c.fail(CloseProtocolError, "unknown opcode")
		}
	}
}