DriverVault/API Server/internal/ocpp/websocket.go

// Package ocpp implements the DriverVault OCPP 1.6J Central System (CSMS) that
// lets the server control an EV charger — start/stop a session, cap the current,
// change availability, reset, and so on — over the persistent WebSocket the
// charger dials out to. It is a separate long-lived subsystem from the read-only
// Anker Solix cloud plugin (internal/plugins/builtin/ankersolix), which can only
// monitor.
//
// The whole API server is dependency-free (stdlib only), so the WebSocket layer
// here is a hand-rolled RFC 6455 implementation rather than gorilla/websocket. It
// is deliberately minimal: text-message oriented (OCPP frames are JSON text),
// single reader / serialized writer, with ping/pong and close handled inline.
package ocpp

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

// wsGUID is the RFC 6455 magic value appended to Sec-WebSocket-Key when deriving
// the Sec-WebSocket-Accept response.
const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// maxMessageBytes caps a single assembled message; OCPP payloads are small, so
// this is a generous ceiling that also bounds a hostile peer.
const maxMessageBytes = 1 << 20 // 1 MiB

// WebSocket opcodes (RFC 6455 §5.2).
const (
	opContinuation = 0x0
	opText         = 0x1
	opBinary       = 0x2
	opClose        = 0x8
	opPing         = 0x9
	opPong         = 0xA
)

// ErrClosed is returned by ReadMessage once the peer has sent a close frame.
var ErrClosed = errors.New("ocpp: connection closed")

// Conn is a minimal RFC 6455 connection, usable as either the server (CSMS) or
// client (proxy→upstream, and the test charge point) end. The only difference is
// masking: per the spec a client masks the frames it sends, a server does not.
type Conn struct {
	raw      net.Conn
	br       *bufio.Reader
	isServer bool

	wmu sync.Mutex // serializes all writes (data + control frames)

	closeOnce sync.Once
	closed    chan struct{}
}

// Upgrade performs the server-side handshake on a hijackable ResponseWriter and
// returns a ready Conn. On success the caller owns the connection and must not
// touch w again. It negotiates the "ocpp1.6" subprotocol when offered.
func Upgrade(w http.ResponseWriter, r *http.Request) (*Conn, error) {
	if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
		!headerHasToken(r.Header.Get("Connection"), "upgrade") {
		return nil, errors.New("ocpp: not a websocket upgrade request")
	}
	key := r.Header.Get("Sec-WebSocket-Key")
	if key == "" {
		return nil, errors.New("ocpp: missing Sec-WebSocket-Key")
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		return nil, errors.New("ocpp: response writer does not support hijacking")
	}
	conn, brw, err := hj.Hijack()
	if err != nil {
		return nil, err
	}

	var subproto string
	for _, p := range splitTokens(r.Header.Get("Sec-WebSocket-Protocol")) {
		if strings.EqualFold(p, "ocpp1.6") {
			subproto = "ocpp1.6"
			break
		}
	}

	var b strings.Builder
	b.WriteString("HTTP/1.1 101 Switching Protocols\r\n")
	b.WriteString("Upgrade: websocket\r\n")
	b.WriteString("Connection: Upgrade\r\n")
	b.WriteString("Sec-WebSocket-Accept: " + acceptKey(key) + "\r\n")
	if subproto != "" {
		b.WriteString("Sec-WebSocket-Protocol: " + subproto + "\r\n")
	}
	b.WriteString("\r\n")
	if _, err := conn.Write([]byte(b.String())); err != nil {
		_ = conn.Close()
		return nil, err
	}
	return &Conn{raw: conn, br: brw.Reader, isServer: true, closed: make(chan struct{})}, nil
}

// Dial opens a client WebSocket to rawURL (ws:// or wss://), offering the given
// subprotocols and any extra request headers (e.g. Authorization for OCPP
// security profile 1). It is used for proxy→upstream links and by tests.
func Dial(ctx context.Context, rawURL string, subprotocols []string, header http.Header) (*Conn, error) {
	u, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("ocpp: bad url: %w", err)
	}
	var (
		secure bool
		port   = u.Port()
	)
	switch strings.ToLower(u.Scheme) {
	case "ws", "http":
		if port == "" {
			port = "80"
		}
	case "wss", "https":
		secure, port = true, orDefault(port, "443")
	default:
		return nil, fmt.Errorf("ocpp: unsupported scheme %q", u.Scheme)
	}

	d := &net.Dialer{}
	raw, err := d.DialContext(ctx, "tcp", net.JoinHostPort(u.Hostname(), port))
	if err != nil {
		return nil, err
	}
	if secure {
		tconn := tls.Client(raw, &tls.Config{ServerName: u.Hostname()})
		if err := tconn.HandshakeContext(ctx); err != nil {
			_ = raw.Close()
			return nil, err
		}
		raw = tconn
	}
	if dl, ok := ctx.Deadline(); ok {
		_ = raw.SetDeadline(dl)
	}

	var keyBytes [16]byte
	if _, err := rand.Read(keyBytes[:]); err != nil {
		_ = raw.Close()
		return nil, err
	}
	key := base64.StdEncoding.EncodeToString(keyBytes[:])

	reqPath := u.RequestURI()
	if reqPath == "" {
		reqPath = "/"
	}
	var b strings.Builder
	b.WriteString("GET " + reqPath + " 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-Version: 13\r\n")
	b.WriteString("Sec-WebSocket-Key: " + key + "\r\n")
	if len(subprotocols) > 0 {
		b.WriteString("Sec-WebSocket-Protocol: " + strings.Join(subprotocols, ", ") + "\r\n")
	}
	for k, vs := range header {
		for _, v := range vs {
			b.WriteString(k + ": " + v + "\r\n")
		}
	}
	b.WriteString("\r\n")
	if _, err := raw.Write([]byte(b.String())); err != nil {
		_ = raw.Close()
		return nil, err
	}

	br := bufio.NewReader(raw)
	resp, err := http.ReadResponse(br, &http.Request{Method: http.MethodGet})
	if err != nil {
		_ = raw.Close()
		return nil, err
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusSwitchingProtocols {
		_ = raw.Close()
		return nil, fmt.Errorf("ocpp: upstream refused upgrade (HTTP %d)", resp.StatusCode)
	}
	if got := resp.Header.Get("Sec-WebSocket-Accept"); got != acceptKey(key) {
		_ = raw.Close()
		return nil, errors.New("ocpp: bad Sec-WebSocket-Accept from upstream")
	}
	// Clear the dial deadline; per-operation deadlines are set explicitly later.
	_ = raw.SetDeadline(time.Time{})
	return &Conn{raw: raw, br: br, isServer: false, closed: make(chan struct{})}, nil
}

// ReadMessage returns the next complete text/binary message, transparently
// answering ping frames and honoring a close frame (after which it returns
// ErrClosed). Control frames never surface to the caller.
func (c *Conn) ReadMessage() ([]byte, error) {
	var (
		buf     []byte
		started bool
	)
	for {
		f, err := c.readFrame()
		if err != nil {
			return nil, err
		}
		switch f.opcode {
		case opPing:
			_ = c.writeFrame(opPong, f.data)
			continue
		case opPong:
			continue
		case opClose:
			_ = c.writeFrame(opClose, f.data)
			c.markClosed()
			return nil, ErrClosed
		case opText, opBinary:
			if started {
				return nil, errors.New("ocpp: new data frame before previous finished")
			}
			started = true
			buf = append(buf, f.data...)
		case opContinuation:
			if !started {
				return nil, errors.New("ocpp: continuation frame with no start")
			}
			buf = append(buf, f.data...)
		default:
			return nil, fmt.Errorf("ocpp: unexpected opcode 0x%x", f.opcode)
		}
		if len(buf) > maxMessageBytes {
			return nil, errors.New("ocpp: message too large")
		}
		if f.fin {
			return buf, nil
		}
	}
}

// WriteMessage sends one OCPP text frame.
func (c *Conn) WriteMessage(data []byte) error {
	select {
	case <-c.closed:
		return ErrClosed
	default:
	}
	return c.writeFrame(opText, data)
}

// SetReadDeadline bounds the next read (used to detect a dead peer between
// heartbeats). A zero time clears it.
func (c *Conn) SetReadDeadline(t time.Time) error { return c.raw.SetReadDeadline(t) }

// Close sends a close frame (best effort) and shuts the underlying connection.
func (c *Conn) Close() error {
	var err error
	c.closeOnce.Do(func() {
		_ = c.writeFrame(opClose, nil)
		close(c.closed)
		err = c.raw.Close()
	})
	return err
}

// markClosed records that the peer initiated close without double-sending.
func (c *Conn) markClosed() {
	c.closeOnce.Do(func() {
		close(c.closed)
		_ = c.raw.Close()
	})
}

// ---- framing -----------------------------------------------------------------

type frame struct {
	fin    bool
	opcode byte
	data   []byte
}

func (c *Conn) readFrame() (frame, error) {
	var h [2]byte
	if _, err := io.ReadFull(c.br, h[:]); err != nil {
		return frame{}, err
	}
	fin := h[0]&0x80 != 0
	if h[0]&0x70 != 0 {
		return frame{}, errors.New("ocpp: reserved bits set (no extensions supported)")
	}
	opcode := h[0] & 0x0f
	masked := h[1]&0x80 != 0
	length := uint64(h[1] & 0x7f)
	switch length {
	case 126:
		var ext [2]byte
		if _, err := io.ReadFull(c.br, ext[:]); err != nil {
			return frame{}, err
		}
		length = uint64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err := io.ReadFull(c.br, ext[:]); err != nil {
			return frame{}, err
		}
		length = binary.BigEndian.Uint64(ext[:])
	}
	if length > maxMessageBytes {
		return frame{}, errors.New("ocpp: frame too large")
	}
	// RFC 6455: the server must receive masked frames; a client must receive
	// unmasked ones.
	if c.isServer && !masked {
		return frame{}, errors.New("ocpp: unmasked frame from client")
	}
	if !c.isServer && masked {
		return frame{}, errors.New("ocpp: masked frame from server")
	}
	var maskKey [4]byte
	if masked {
		if _, err := io.ReadFull(c.br, maskKey[:]); err != nil {
			return frame{}, err
		}
	}
	data := make([]byte, length)
	if _, err := io.ReadFull(c.br, data); err != nil {
		return frame{}, err
	}
	if masked {
		for i := range data {
			data[i] ^= maskKey[i%4]
		}
	}
	return frame{fin: fin, opcode: opcode, data: data}, nil
}

func (c *Conn) writeFrame(opcode byte, data []byte) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()

	mask := !c.isServer // clients mask their frames
	var head [14]byte
	head[0] = 0x80 | opcode // FIN + opcode (we never fragment outgoing messages)
	n := 2
	l := len(data)
	switch {
	case l <= 125:
		head[1] = byte(l)
	case l <= 0xFFFF:
		head[1] = 126
		binary.BigEndian.PutUint16(head[2:4], uint16(l))
		n = 4
	default:
		head[1] = 127
		binary.BigEndian.PutUint64(head[2:10], uint64(l))
		n = 10
	}

	var maskKey [4]byte
	if mask {
		head[1] |= 0x80
		if _, err := rand.Read(maskKey[:]); err != nil {
			return err
		}
		copy(head[n:n+4], maskKey[:])
		n += 4
	}

	_ = c.raw.SetWriteDeadline(time.Now().Add(30 * time.Second))
	defer c.raw.SetWriteDeadline(time.Time{})

	if _, err := c.raw.Write(head[:n]); err != nil {
		return err
	}
	if l == 0 {
		return nil
	}
	if mask {
		masked := make([]byte, l)
		for i := 0; i < l; i++ {
			masked[i] = data[i] ^ maskKey[i%4]
		}
		_, err := c.raw.Write(masked)
		return err
	}
	_, err := c.raw.Write(data)
	return err
}

// ---- small helpers -----------------------------------------------------------

func acceptKey(key string) string {
	h := sha1.New()
	_, _ = io.WriteString(h, key+wsGUID)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

// headerHasToken reports whether a comma-separated header value contains token
// (case-insensitive) — e.g. Connection: keep-alive, Upgrade.
func headerHasToken(value, token string) bool {
	for _, t := range splitTokens(value) {
		if strings.EqualFold(t, token) {
			return true
		}
	}
	return false
}

func splitTokens(value string) []string {
	if value == "" {
		return nil
	}
	parts := strings.Split(value, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}

func orDefault(v, def string) string {
	if v == "" {
		return def
	}
	return v
}