DriverVault/API Server/internal/mqtt/client.go

// Package mqtt is a minimal MQTT 3.1.1 client, written against the OASIS spec
// rather than pulled in as a dependency: the whole API server is stdlib-only, so
// this sits beside internal/ocpp's hand-rolled RFC 6455 WebSocket for the same
// reason.
//
// It is deliberately scoped to what the Greencell EVSE connector needs
// (internal/plugins/builtin/greencell): connect to a broker, subscribe to a
// handful of topics, publish a discovery broadcast, and read the telemetry that
// comes back. Concretely that means:
//
//   - QoS 0 for everything this client *sends* — subscriptions request QoS 0 and
//     publishes are QoS 0. The Greencell device and Home Assistant's own
//     integration both work at QoS 0.
//   - Inbound PUBLISH at QoS 0, 1 and 2 is accepted and acknowledged correctly,
//     because the QoS of a delivery is the broker's choice, not ours.
//   - Clean session only. There is no persistent session, no offline queue and no
//     automatic reconnect: a connection lives for the length of one plugin call
//     and is closed again, which is exactly how the plugin manager constructs and
//     tears down a plugin instance per request.
//
// A Client is safe for concurrent use; writes are serialized and a single reader
// goroutine owns the connection.
package mqtt

import (
	"bufio"
	"context"
	"crypto/rand"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"
)

// Defaults applied by Options.normalize when a field is left zero.
const (
	defaultKeepalive      = 30 * time.Second
	defaultConnectTimeout = 10 * time.Second
	defaultBuffer         = 256
)

// ErrClosed is returned by Publish and Subscribe once the connection is gone.
var ErrClosed = errors.New("mqtt: connection closed")

// Options configures a broker connection.
type Options struct {
	// Address is the broker's host:port. Required.
	Address string
	// TLS wraps the connection in TLS. TLSConfig overrides the default, which
	// verifies the broker against the system roots using the host from Address.
	TLS       bool
	TLSConfig *tls.Config
	// ClientID identifies this session to the broker. Empty means a random one.
	ClientID string
	// Username and Password are sent in CONNECT when Username is non-empty.
	Username string
	Password string
	// Keepalive is the interval the broker is told to expect traffic within; the
	// client pings at half of it.
	Keepalive time.Duration
	// ConnectTimeout bounds the TCP/TLS dial and the wait for CONNACK.
	ConnectTimeout time.Duration
	// Buffer sizes the delivery channel returned by Messages.
	Buffer int
}

func (o *Options) normalize() {
	if o.Keepalive <= 0 {
		o.Keepalive = defaultKeepalive
	}
	if o.ConnectTimeout <= 0 {
		o.ConnectTimeout = defaultConnectTimeout
	}
	if o.Buffer <= 0 {
		o.Buffer = defaultBuffer
	}
	if o.ClientID == "" {
		o.ClientID = randomClientID()
	}
}

// Message is one delivered PUBLISH.
type Message struct {
	Topic   string
	Payload []byte
	Retain  bool
}

// Client is a connected MQTT session.
type Client struct {
	raw net.Conn
	br  *bufio.Reader

	keepalive time.Duration

	wmu sync.Mutex // serializes writes; the spec requires whole packets

	msgs chan Message

	mu      sync.Mutex
	nextID  uint16
	pending map[uint16]chan []byte // packet id -> acknowledgement body
	err     error

	closeOnce sync.Once
	closed    chan struct{}
}

// Connect dials the broker and completes the MQTT handshake. The returned Client
// owns the connection; call Close when done.
func Connect(ctx context.Context, opt Options) (*Client, error) {
	opt.normalize()
	if opt.Address == "" {
		return nil, errors.New("mqtt: broker address is required")
	}

	dialer := &net.Dialer{Timeout: opt.ConnectTimeout}
	var (
		conn net.Conn
		err  error
	)
	if opt.TLS {
		cfg := opt.TLSConfig
		if cfg == nil {
			host, _, splitErr := net.SplitHostPort(opt.Address)
			if splitErr != nil {
				host = opt.Address
			}
			cfg = &tls.Config{ServerName: host, MinVersion: tls.VersionTLS12}
		}
		conn, err = (&tls.Dialer{NetDialer: dialer, Config: cfg}).DialContext(ctx, "tcp", opt.Address)
	} else {
		conn, err = dialer.DialContext(ctx, "tcp", opt.Address)
	}
	if err != nil {
		return nil, fmt.Errorf("mqtt: cannot reach broker %s: %w", opt.Address, err)
	}

	c := &Client{
		raw:       conn,
		br:        bufio.NewReader(conn),
		keepalive: opt.Keepalive,
		msgs:      make(chan Message, opt.Buffer),
		nextID:    1,
		pending:   map[uint16]chan []byte{},
		closed:    make(chan struct{}),
	}

	deadline := time.Now().Add(opt.ConnectTimeout)
	if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
		deadline = d
	}
	if err := c.handshake(opt, deadline); err != nil {
		_ = conn.Close()
		return nil, err
	}

	go c.readLoop()
	go c.pingLoop()
	return c, nil
}

// handshake writes CONNECT and waits for an accepting CONNACK. It runs before the
// read loop starts, so it reads the response itself.
func (c *Client) handshake(opt Options, deadline time.Time) error {
	if len(opt.ClientID) > 65535 || len(opt.Username) > 65535 || len(opt.Password) > 65535 {
		return errors.New("mqtt: client id or credentials too long")
	}

	var flags byte = 0x02 // clean session
	if opt.Username != "" {
		flags |= 0x80
		if opt.Password != "" {
			flags |= 0x40
		}
	}

	body := make([]byte, 0, 64)
	body = encodeString(body, protocolName)
	body = append(body, protocolLevel, flags)
	body = encodeUint16(body, uint16(opt.Keepalive/time.Second))
	body = encodeString(body, opt.ClientID)
	if opt.Username != "" {
		body = encodeString(body, opt.Username)
		if opt.Password != "" {
			body = encodeString(body, opt.Password)
		}
	}

	if err := c.raw.SetWriteDeadline(deadline); err != nil {
		return err
	}
	if _, err := c.raw.Write(buildPacket(pktConnect, 0, body)); err != nil {
		return fmt.Errorf("mqtt: sending CONNECT: %w", err)
	}
	if err := c.raw.SetWriteDeadline(time.Time{}); err != nil {
		return err
	}

	if err := c.raw.SetReadDeadline(deadline); err != nil {
		return err
	}
	pkt, err := readPacket(c.br)
	if err != nil {
		return fmt.Errorf("mqtt: waiting for CONNACK: %w", err)
	}
	if err := c.raw.SetReadDeadline(time.Time{}); err != nil {
		return err
	}
	if pkt.typ != pktConnack || len(pkt.body) < 2 {
		return errors.New("mqtt: broker did not answer with a CONNACK")
	}
	return connackError(pkt.body[1])
}

// Messages returns the stream of delivered publications. It is closed when the
// connection ends; check Err for the reason.
func (c *Client) Messages() <-chan Message { return c.msgs }

// Err reports why the connection ended, or nil while it is healthy or after a
// clean Close.
func (c *Client) Err() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.err
}

// Subscribe requests QoS 0 on every given topic filter in one SUBSCRIBE and waits
// for the broker's SUBACK.
func (c *Client) Subscribe(ctx context.Context, topics ...string) error {
	if len(topics) == 0 {
		return nil
	}
	id, ack := c.reserveID()
	defer c.releaseID(id)

	body := encodeUint16(make([]byte, 0, 16*len(topics)), id)
	for _, t := range topics {
		if t == "" || len(t) > 65535 {
			return fmt.Errorf("mqtt: invalid topic filter %q", t)
		}
		body = encodeString(body, t)
		body = append(body, 0) // requested QoS 0
	}
	// SUBSCRIBE reserves the fixed-header flag bits 0010 (spec 3.8.1).
	if err := c.write(buildPacket(pktSubscribe, 0x02, body)); err != nil {
		return err
	}

	select {
	case codes := <-ack:
		for i, code := range codes {
			if code == 0x80 {
				return fmt.Errorf("mqtt: broker refused subscription to %s", topics[min(i, len(topics)-1)])
			}
		}
		return nil
	case <-c.closed:
		return c.closedErr()
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Publish sends a QoS 0, non-retained publication. QoS 0 is fire-and-forget, so
// it returns as soon as the packet is on the wire.
func (c *Client) Publish(ctx context.Context, topic string, payload []byte) error {
	if topic == "" || len(topic) > 65535 {
		return fmt.Errorf("mqtt: invalid topic %q", topic)
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	body := encodeString(make([]byte, 0, len(topic)+len(payload)+2), topic)
	body = append(body, payload...)
	return c.write(buildPacket(pktPublish, 0, body))
}

// Close ends the session, sending DISCONNECT while the connection still works.
func (c *Client) Close() error {
	c.closeOnce.Do(func() {
		// Best effort: a broker that has already gone away needs no goodbye.
		_ = c.write(buildPacket(pktDisconnect, 0, nil))
		close(c.closed)
		_ = c.raw.Close()
	})
	return nil
}

// write sends one whole packet under the write lock.
func (c *Client) write(b []byte) error {
	select {
	case <-c.closed:
		return c.closedErr()
	default:
	}
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if err := c.raw.SetWriteDeadline(time.Now().Add(c.keepalive)); err != nil {
		return err
	}
	_, err := c.raw.Write(b)
	return err
}

// closedErr reports the failure that ended the connection, falling back to
// ErrClosed after a clean shutdown.
func (c *Client) closedErr() error {
	if err := c.Err(); err != nil {
		return err
	}
	return ErrClosed
}

// reserveID allocates a packet identifier and the channel its acknowledgement
// will be delivered on. Identifier 0 is not valid on the wire, so it is skipped
// on wrap.
func (c *Client) reserveID() (uint16, chan []byte) {
	ch := make(chan []byte, 1)
	c.mu.Lock()
	defer c.mu.Unlock()
	id := c.nextID
	c.nextID++
	if c.nextID == 0 {
		c.nextID = 1
	}
	c.pending[id] = ch
	return id, ch
}

func (c *Client) releaseID(id uint16) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.pending, id)
}

// resolve hands an acknowledgement body to whoever is waiting on that packet id.
func (c *Client) resolve(id uint16, body []byte) {
	c.mu.Lock()
	ch, ok := c.pending[id]
	c.mu.Unlock()
	if !ok {
		return
	}
	select {
	case ch <- body:
	default: // a waiter that already gave up
	}
}

// fail records the error that ended the connection and tears it down.
func (c *Client) fail(err error) {
	c.mu.Lock()
	if c.err == nil {
		c.err = err
	}
	c.mu.Unlock()
	_ = c.Close()
}

// readLoop owns the read side of the connection for the life of the session.
func (c *Client) readLoop() {
	defer close(c.msgs)
	for {
		// The broker is expected to answer our pings, so silence for two
		// keep-alive periods means the connection is dead.
		if err := c.raw.SetReadDeadline(time.Now().Add(2 * c.keepalive)); err != nil {
			c.fail(err)
			return
		}
		pkt, err := readPacket(c.br)
		if err != nil {
			select {
			case <-c.closed:
				return // our own Close raced the read; not a failure
			default:
			}
			c.fail(fmt.Errorf("mqtt: reading from broker: %w", err))
			return
		}
		if !c.dispatch(pkt) {
			return
		}
	}
}

// dispatch handles one received packet and reports whether the loop continues.
func (c *Client) dispatch(pkt packet) bool {
	switch pkt.typ {
	case pktPublish:
		return c.deliver(pkt)
	case pktSuback, pktUnsuback, pktPuback:
		if len(pkt.body) >= 2 {
			c.resolve(binary.BigEndian.Uint16(pkt.body), pkt.body[2:])
		}
	case pktPubrel:
		// Second half of an inbound QoS 2 delivery: the broker released the
		// message, so complete the exchange.
		if len(pkt.body) >= 2 {
			body := encodeUint16(nil, binary.BigEndian.Uint16(pkt.body))
			if err := c.write(buildPacket(pktPubcomp, 0, body)); err != nil {
				c.fail(err)
				return false
			}
		}
	case pktPingresp, pktPubrec, pktPubcomp:
		// Nothing to do: we send no QoS >0 publications of our own.
	}
	return true
}

// deliver parses an inbound PUBLISH, acknowledges it at its QoS, and hands it to
// the consumer.
func (c *Client) deliver(pkt packet) bool {
	qos := (pkt.flags >> 1) & 0x03
	topic, rest, err := readString(pkt.body)
	if err != nil {
		c.fail(fmt.Errorf("mqtt: malformed PUBLISH: %w", err))
		return false
	}
	if qos > 0 {
		if len(rest) < 2 {
			c.fail(errors.New("mqtt: PUBLISH missing packet identifier"))
			return false
		}
		id := binary.BigEndian.Uint16(rest)
		rest = rest[2:]
		// QoS 1 completes with PUBACK; QoS 2 starts the four-way exchange whose
		// PUBREL half is handled in dispatch.
		ackType := pktPuback
		if qos == 2 {
			ackType = pktPubrec
		}
		if err := c.write(buildPacket(ackType, 0, encodeUint16(nil, id))); err != nil {
			c.fail(err)
			return false
		}
	}

	msg := Message{Topic: topic, Payload: rest, Retain: pkt.flags&0x01 != 0}
	select {
	case c.msgs <- msg:
	case <-c.closed:
		return false
	}
	return true
}

// pingLoop keeps the session alive, pinging at half the negotiated interval so a
// single lost PINGREQ does not expire it.
func (c *Client) pingLoop() {
	t := time.NewTicker(c.keepalive / 2)
	defer t.Stop()
	for {
		select {
		case <-c.closed:
			return
		case <-t.C:
			if err := c.write(buildPacket(pktPingreq, 0, nil)); err != nil {
				c.fail(err)
				return
			}
		}
	}
}

// randomClientID builds an identifier the broker has not seen before, so two
// concurrent plugin calls never evict each other's session.
func randomClientID() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return fmt.Sprintf("drivervault-%d", time.Now().UnixNano())
	}
	return "drivervault-" + hex.EncodeToString(b[:])
}