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

package mqtt

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"net"
	"strings"
	"sync"
	"testing"
	"time"
)

// ---- a broker just real enough to test the client against ---------------------

// brokerConn is the server side of one connection, speaking the same wire codec
// the client does.
type brokerConn struct {
	conn net.Conn
	br   *bufio.Reader
}

func (b *brokerConn) read() (packet, error) { return readPacket(b.br) }

func (b *brokerConn) send(typ, flags byte, body []byte) error {
	_, err := b.conn.Write(buildPacket(typ, flags, body))
	return err
}

// accept reads the CONNECT and answers with an accepting CONNACK, returning the
// CONNECT body so a test can assert on what the client sent.
func (b *brokerConn) accept(t *testing.T) []byte {
	t.Helper()
	pkt, err := b.read()
	if err != nil {
		t.Fatalf("reading CONNECT: %v", err)
	}
	if pkt.typ != pktConnect {
		t.Fatalf("first packet type = %d, want CONNECT(%d)", pkt.typ, pktConnect)
	}
	if err := b.send(pktConnack, 0, []byte{0, 0}); err != nil {
		t.Fatalf("sending CONNACK: %v", err)
	}
	return pkt.body
}

// publish pushes a PUBLISH to the client at the given QoS.
func (b *brokerConn) publish(topic string, payload []byte, qos byte, id uint16, retain bool) error {
	body := encodeString(nil, topic)
	if qos > 0 {
		body = encodeUint16(body, id)
	}
	body = append(body, payload...)
	flags := qos << 1
	if retain {
		flags |= 0x01
	}
	return b.send(pktPublish, flags, body)
}

// startBroker listens on a loopback port and hands the first connection to
// handle. It returns the address to dial.
func startBroker(t *testing.T, handle func(b *brokerConn)) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	var wg sync.WaitGroup
	// Close the listener before waiting: cleanups run last-registered first, so
	// registering both together keeps the order right.
	t.Cleanup(func() { _ = ln.Close(); wg.Wait() })
	wg.Add(1)
	go func() {
		defer wg.Done()
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		handle(&brokerConn{conn: conn, br: bufio.NewReader(conn)})
	}()
	return ln.Addr().String()
}

func dial(t *testing.T, addr string, opt Options) *Client {
	t.Helper()
	opt.Address = addr
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	c, err := Connect(ctx, opt)
	if err != nil {
		t.Fatalf("Connect: %v", err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

// ---- wire codec --------------------------------------------------------------

func TestRemainingLengthRoundTrip(t *testing.T) {
	// The boundaries of each varint width, per the MQTT spec's own table.
	for _, n := range []int{0, 1, 127, 128, 16383, 16384, 2097151, 2097152, maxRemainingLength} {
		encoded := encodeRemainingLength(nil, n)
		got, err := readRemainingLength(bufio.NewReader(strings.NewReader(string(encoded))))
		if err != nil {
			t.Fatalf("readRemainingLength(%d): %v", n, err)
		}
		if got != n {
			t.Errorf("round trip %d -> %d", n, got)
		}
	}
}

func TestRemainingLengthRejectsOverlongVarint(t *testing.T) {
	// Five continuation bytes is malformed; the spec caps the field at four.
	br := bufio.NewReader(strings.NewReader("\xff\xff\xff\xff\xff"))
	if _, err := readRemainingLength(br); err == nil {
		t.Fatal("expected an error for a five-byte remaining length")
	}
}

func TestReadString(t *testing.T) {
	b := encodeString(nil, "/greencell/broadcast")
	b = append(b, 'x', 'y')
	s, rest, err := readString(b)
	if err != nil {
		t.Fatalf("readString: %v", err)
	}
	if s != "/greencell/broadcast" {
		t.Errorf("string = %q", s)
	}
	if string(rest) != "xy" {
		t.Errorf("rest = %q, want xy", rest)
	}

	if _, _, err := readString([]byte{0x00, 0x05, 'a'}); err == nil {
		t.Error("expected an error for a truncated string")
	}
	if _, _, err := readString([]byte{0x00}); err == nil {
		t.Error("expected an error for a truncated length")
	}
}

func TestConnackError(t *testing.T) {
	if err := connackError(0); err != nil {
		t.Errorf("code 0 should be accepted, got %v", err)
	}
	if err := connackError(4); err == nil || !strings.Contains(err.Error(), "username or password") {
		t.Errorf("code 4 = %v, want a credentials error", err)
	}
	if err := connackError(9); err == nil {
		t.Error("an unknown code should still be an error")
	}
}

// ---- handshake ---------------------------------------------------------------

func TestConnectSendsCredentialsAndKeepalive(t *testing.T) {
	var got []byte
	done := make(chan struct{})
	addr := startBroker(t, func(b *brokerConn) {
		got = b.accept(t)
		close(done)
		// Hold the connection open so the client stays up for the assertions.
		_, _ = b.read()
	})

	dial(t, addr, Options{ClientID: "dv-test", Username: "user", Password: "pass", Keepalive: 42 * time.Second})
	<-done

	name, rest, err := readString(got)
	if err != nil || name != protocolName {
		t.Fatalf("protocol name = %q (err %v), want MQTT", name, err)
	}
	if rest[0] != protocolLevel {
		t.Errorf("protocol level = %d, want %d", rest[0], protocolLevel)
	}
	flags := rest[1]
	if flags&0x02 == 0 {
		t.Error("clean session flag should be set")
	}
	if flags&0x80 == 0 || flags&0x40 == 0 {
		t.Errorf("username+password flags should be set, got %#x", flags)
	}
	if ka := binary.BigEndian.Uint16(rest[2:]); ka != 42 {
		t.Errorf("keepalive = %d, want 42", ka)
	}

	id, rest, err := readString(rest[4:])
	if err != nil || id != "dv-test" {
		t.Fatalf("client id = %q (err %v)", id, err)
	}
	user, rest, err := readString(rest)
	if err != nil || user != "user" {
		t.Fatalf("username = %q (err %v)", user, err)
	}
	pass, _, err := readString(rest)
	if err != nil || pass != "pass" {
		t.Fatalf("password = %q (err %v)", pass, err)
	}
}

func TestConnectOmitsCredentialsWhenUnset(t *testing.T) {
	var got []byte
	done := make(chan struct{})
	addr := startBroker(t, func(b *brokerConn) {
		got = b.accept(t)
		close(done)
		_, _ = b.read()
	})

	dial(t, addr, Options{ClientID: "dv-test"})
	<-done

	_, rest, err := readString(got)
	if err != nil {
		t.Fatalf("readString: %v", err)
	}
	if flags := rest[1]; flags&0xC0 != 0 {
		t.Errorf("credential flags should be clear on an anonymous connect, got %#x", flags)
	}
}

func TestConnectRefused(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		if _, err := b.read(); err != nil {
			return
		}
		_ = b.send(pktConnack, 0, []byte{0, 5}) // not authorized
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err := Connect(ctx, Options{Address: addr})
	if err == nil {
		t.Fatal("expected the refusal to surface as an error")
	}
	if !strings.Contains(err.Error(), "not authorized") {
		t.Errorf("error = %v, want the CONNACK reason", err)
	}
}

func TestConnectRejectsNonConnack(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		if _, err := b.read(); err != nil {
			return
		}
		_ = b.send(pktPingresp, 0, nil)
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if _, err := Connect(ctx, Options{Address: addr}); err == nil {
		t.Fatal("expected an error when the broker answers something other than CONNACK")
	}
}

func TestConnectUnreachableBroker(t *testing.T) {
	// Port 0 never accepts, so this fails at dial rather than at handshake.
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err := Connect(ctx, Options{Address: "127.0.0.1:0", ConnectTimeout: time.Second})
	if err == nil {
		t.Fatal("expected a dial error")
	}
	if !strings.Contains(err.Error(), "cannot reach broker") {
		t.Errorf("error = %v, want a broker-unreachable message", err)
	}
}

func TestConnectRequiresAddress(t *testing.T) {
	if _, err := Connect(context.Background(), Options{}); err == nil {
		t.Fatal("expected an error when no broker address is configured")
	}
}

// ---- subscribe / publish -----------------------------------------------------

func TestSubscribeAndDeliver(t *testing.T) {
	var topics []string
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		pkt, err := b.read()
		if err != nil {
			return
		}
		if pkt.typ != pktSubscribe {
			t.Errorf("packet type = %d, want SUBSCRIBE(%d)", pkt.typ, pktSubscribe)
			return
		}
		if pkt.flags != 0x02 {
			t.Errorf("SUBSCRIBE flags = %#x, want 0x02", pkt.flags)
		}
		id := binary.BigEndian.Uint16(pkt.body)
		rest := pkt.body[2:]
		for len(rest) > 0 {
			var topic string
			topic, rest, err = readString(rest)
			if err != nil {
				return
			}
			topics = append(topics, topic)
			rest = rest[1:] // requested QoS
		}
		_ = b.send(pktSuback, 0, append(encodeUint16(nil, id), 0, 0))
		_ = b.publish("/greencell/evse/SN1/power", []byte(`{"momentary":7360}`), 0, 0, true)
		_, _ = b.read()
	})

	c := dial(t, addr, Options{})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := c.Subscribe(ctx, "/greencell/broadcast/device", "/greencell/evse/SN1/power"); err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	if len(topics) != 2 || topics[0] != "/greencell/broadcast/device" {
		t.Errorf("broker saw topics %v", topics)
	}

	select {
	case msg := <-c.Messages():
		if msg.Topic != "/greencell/evse/SN1/power" {
			t.Errorf("topic = %q", msg.Topic)
		}
		if string(msg.Payload) != `{"momentary":7360}` {
			t.Errorf("payload = %q", msg.Payload)
		}
		if !msg.Retain {
			t.Error("retain flag should have been carried through")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the delivery")
	}
}

func TestSubscribeRefused(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		pkt, err := b.read()
		if err != nil {
			return
		}
		id := binary.BigEndian.Uint16(pkt.body)
		_ = b.send(pktSuback, 0, append(encodeUint16(nil, id), 0x80))
		_, _ = b.read()
	})

	c := dial(t, addr, Options{})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := c.Subscribe(ctx, "/greencell/broadcast/device")
	if err == nil || !strings.Contains(err.Error(), "refused subscription") {
		t.Fatalf("Subscribe error = %v, want a refusal", err)
	}
}

func TestSubscribeNoTopicsIsNoOp(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		_, _ = b.read()
	})
	c := dial(t, addr, Options{})
	if err := c.Subscribe(context.Background()); err != nil {
		t.Fatalf("Subscribe with no topics = %v, want nil", err)
	}
}

func TestPublish(t *testing.T) {
	type sent struct {
		topic   string
		payload string
	}
	got := make(chan sent, 1)
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		pkt, err := b.read()
		if err != nil {
			return
		}
		if pkt.typ != pktPublish {
			t.Errorf("packet type = %d, want PUBLISH(%d)", pkt.typ, pktPublish)
			return
		}
		if qos := (pkt.flags >> 1) & 0x03; qos != 0 {
			t.Errorf("published QoS = %d, want 0", qos)
		}
		topic, rest, err := readString(pkt.body)
		if err != nil {
			return
		}
		got <- sent{topic, string(rest)}
		_, _ = b.read()
	})

	c := dial(t, addr, Options{})
	if err := c.Publish(context.Background(), "/greencell/broadcast", []byte(`{"name":"BROADCAST"}`)); err != nil {
		t.Fatalf("Publish: %v", err)
	}
	select {
	case s := <-got:
		if s.topic != "/greencell/broadcast" || s.payload != `{"name":"BROADCAST"}` {
			t.Errorf("broker received %+v", s)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the publish")
	}
}

func TestPublishRejectsEmptyTopic(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		_, _ = b.read()
	})
	c := dial(t, addr, Options{})
	if err := c.Publish(context.Background(), "", nil); err == nil {
		t.Fatal("expected an error for an empty topic")
	}
}

// TestInboundQoS1IsAcknowledged covers the case the client does not choose: the
// broker may upgrade a delivery to QoS 1, and an unacknowledged one is redelivered
// forever.
func TestInboundQoS1IsAcknowledged(t *testing.T) {
	acked := make(chan uint16, 1)
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		_ = b.publish("/greencell/evse/SN1/status", []byte(`{"state":"CHARGING"}`), 1, 77, false)
		for {
			pkt, err := b.read()
			if err != nil {
				return
			}
			if pkt.typ == pktPuback && len(pkt.body) >= 2 {
				acked <- binary.BigEndian.Uint16(pkt.body)
				return
			}
		}
	})

	c := dial(t, addr, Options{})
	select {
	case msg := <-c.Messages():
		if string(msg.Payload) != `{"state":"CHARGING"}` {
			t.Errorf("payload = %q — the packet identifier should be stripped", msg.Payload)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the QoS 1 delivery")
	}
	select {
	case id := <-acked:
		if id != 77 {
			t.Errorf("PUBACK id = %d, want 77", id)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the PUBACK")
	}
}

// TestInboundQoS2Completes walks the four-way exchange the broker drives.
func TestInboundQoS2Completes(t *testing.T) {
	done := make(chan byte, 2)
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		_ = b.publish("/greencell/evse/SN1/status", []byte(`{"state":"IDLE"}`), 2, 88, false)
		for {
			pkt, err := b.read()
			if err != nil {
				return
			}
			switch pkt.typ {
			case pktPubrec:
				done <- pktPubrec
				_ = b.send(pktPubrel, 0x02, encodeUint16(nil, 88))
			case pktPubcomp:
				done <- pktPubcomp
				return
			}
		}
	})

	c := dial(t, addr, Options{})
	select {
	case <-c.Messages():
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the QoS 2 delivery")
	}
	for _, want := range []byte{pktPubrec, pktPubcomp} {
		select {
		case got := <-done:
			if got != want {
				t.Fatalf("exchange step = %d, want %d", got, want)
			}
		case <-time.After(5 * time.Second):
			t.Fatalf("timed out waiting for step %d", want)
		}
	}
}

// ---- teardown ----------------------------------------------------------------

func TestMessagesCloseAndErrOnBrokerDisconnect(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		_ = b.conn.Close()
	})

	c := dial(t, addr, Options{})
	select {
	case _, ok := <-c.Messages():
		if ok {
			t.Fatal("expected the message channel to close, not deliver")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the channel to close")
	}
	if c.Err() == nil {
		t.Error("Err should report why the connection ended")
	}
}

func TestCloseIsIdempotentAndSilencesErr(t *testing.T) {
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		for {
			if _, err := b.read(); err != nil {
				return
			}
		}
	})

	c := dial(t, addr, Options{})
	if err := c.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}
	if err := c.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}
	if err := c.Err(); err != nil {
		t.Errorf("a clean Close should leave Err nil, got %v", err)
	}
	if err := c.Publish(context.Background(), "/x", nil); !errors.Is(err, ErrClosed) {
		t.Errorf("Publish after Close = %v, want ErrClosed", err)
	}
}

func TestSubscribeHonoursContextCancellation(t *testing.T) {
	// A broker that never answers the SUBSCRIBE: the caller's deadline is what
	// has to end the wait.
	addr := startBroker(t, func(b *brokerConn) {
		b.accept(t)
		for {
			if _, err := b.read(); err != nil {
				return
			}
		}
	})

	c := dial(t, addr, Options{})
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()
	if err := c.Subscribe(ctx, "/greencell/broadcast/device"); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("Subscribe = %v, want a deadline error", err)
	}
}