marionette/

lib.rs

//! Minimal scripted MQTT v5 "broker" for driving protocol edge cases that a
//! real broker will never produce (unsolicited packets, withheld CONNACKs,
//! stalled sockets, exotic CONNACK properties).
//!
//! This is not a conforming broker: each test scripts the exact packets to
//! read and write. Packets are encoded/decoded with `finmqtt-core`, so tests
//! can construct and inspect them with the regular packet types.

#![allow(
    clippy::expect_used,
    clippy::panic,
    reason = "this is a test helper crate"
)]

use std::io::{self, ErrorKind};
use std::time::Duration;

use finmqtt_core::codec::{GrowBuffer, PacketReader};
use finmqtt_core::packet::{Ack, Connack, Packet, ReasonCode, Suback};
use finmqtt_core::property::{AckProperties, ConnackPropertiesBuf, Property, SubackProperties};
use finmqtt_core::wire::{FixedHeader, PacketId, PacketType};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::time::timeout;

/// Default deadline for blocking-style helpers ([`Broker::accept`],
/// [`Conn::read_packet`]).
const DEADLINE: Duration = Duration::from_secs(10);

/// A TCP listener that hands out scripted MQTT connections.
pub struct Broker {
    listener: TcpListener,
}

impl Broker {
    /// Binds to an ephemeral IPv6 port.
    ///
    /// Listening on `[::]` accepts IPv6 clients and, on the usual dual-stack
    /// configuration, IPv4-mapped clients (e.g. `127.0.0.1`) too. On a rare
    /// IPv6-only host, clients simply connect over IPv6.
    ///
    /// # Panics
    ///
    /// Panics if binding fails.
    pub async fn bind() -> Self {
        let listener = TcpListener::bind("[::]:0")
            .await
            .expect("failed to bind fake broker");
        Self { listener }
    }

    /// The TCP port the fake broker is listening on.
    ///
    /// # Panics
    ///
    /// Panics if the listener's local address cannot be queried.
    #[must_use]
    pub fn port(&self) -> u16 {
        self.listener
            .local_addr()
            .expect("failed to get local addr")
            .port()
    }

    /// The loopback host a client should connect to, matching the bound
    /// address family. Pair it with [`port`](Self::port) so test code never
    /// has to assume IPv4 or dual-stack.
    #[must_use]
    pub fn host(&self) -> &'static str {
        if self
            .listener
            .local_addr()
            .map_or(true, |addr| addr.is_ipv6())
        {
            "::1"
        } else {
            "127.0.0.1"
        }
    }

    /// Accepts the next client connection.
    ///
    /// # Panics
    ///
    /// Panics if no client connects within 10 s.
    pub async fn accept(&self) -> Conn {
        self.try_accept(DEADLINE)
            .await
            .expect("timed out waiting for a client connection")
    }

    /// Accepts the next client connection, returning `None` if no client
    /// connects within `dur`.
    ///
    /// # Panics
    ///
    /// Panics on listener I/O errors.
    pub async fn try_accept(&self, dur: Duration) -> Option<Conn> {
        match timeout(dur, self.listener.accept()).await {
            Ok(Ok((stream, _addr))) => Some(Conn {
                stream,
                reader: PacketReader::with_capacity(8 * 1024),
            }),
            Ok(Err(e)) => panic!("fake broker accept failed: {e}"),
            Err(_elapsed) => None,
        }
    }
}

/// A single scripted broker-side connection.
pub struct Conn {
    stream: TcpStream,
    reader: PacketReader<GrowBuffer>,
}

/// A packet received from the client: fixed header plus body bytes.
///
/// Use [`RecvPacket::decode`] to view it as a [`Packet`].
pub struct RecvPacket {
    pub header: FixedHeader,
    pub body: Vec<u8>,
}

impl RecvPacket {
    /// Decodes the packet, panicking on malformed data.
    ///
    /// # Panics
    ///
    /// Panics if the body cannot be decoded for the header's packet type.
    #[must_use]
    pub fn decode(&self) -> Packet<'_> {
        Packet::decode(self.header, &self.body)
            .unwrap_or_else(|e| panic!("client sent malformed {:?}: {e:?}", self.header))
    }

    /// The packet type from the fixed header.
    #[must_use]
    pub const fn packet_type(&self) -> PacketType {
        self.header.packet_type
    }
}

impl Conn {
    /// Reads the next packet from the client.
    ///
    /// # Panics
    ///
    /// Panics if no packet arrives within 10 s or the connection closes.
    pub async fn read_packet(&mut self) -> RecvPacket {
        let result = timeout(DEADLINE, self.read_packet_or_eof())
            .await
            .expect("timed out waiting for a packet from the client");
        result.expect("client closed the connection while a packet was expected")
    }

    /// Reads the next packet, returning `None` on timeout or if the client
    /// closed the connection cleanly.
    ///
    /// # Panics
    ///
    /// Panics on I/O errors or if the connection drops mid-packet.
    pub async fn try_read_packet(&mut self, dur: Duration) -> Option<RecvPacket> {
        timeout(dur, self.read_packet_or_eof())
            .await
            .unwrap_or_default()
    }

    /// Reads one packet; `None` on a clean EOF before the next packet starts.
    async fn read_packet_or_eof(&mut self) -> Option<RecvPacket> {
        let mut chunk = [0u8; 1024];
        loop {
            match self.reader.peek() {
                Ok(Some((header, total))) => {
                    let body = self.reader.body(header, total).to_vec();
                    self.reader.consume(total);
                    return Some(RecvPacket { header, body });
                }
                Ok(None) => {}
                Err(e) => panic!("client sent a malformed packet header: {e:?}"),
            }
            match self.stream.read(&mut chunk).await {
                Ok(0) => {
                    assert!(
                        self.reader.buffered_len() == 0,
                        "connection dropped mid packet"
                    );
                    return None;
                }
                Ok(n) => self
                    .reader
                    .feed(&chunk[..n])
                    .expect("GrowBuffer never reports a capacity limit"),
                Err(e) if e.kind() == ErrorKind::UnexpectedEof => return None,
                Err(e) => panic!("fake broker read failed: {e}"),
            }
        }
    }

    /// Encodes and writes a packet to the client.
    ///
    /// # Panics
    ///
    /// Panics if encoding or the socket write fails.
    pub async fn write_packet(&mut self, pkt: &Packet<'_>) {
        let mut buf = vec![0u8; pkt.encoded_len().expect("packet too large to encode")];
        let n = pkt.encode(&mut buf).expect("failed to encode packet");
        self.write_raw(&buf[..n]).await;
    }

    /// Writes raw bytes to the client, bypassing packet encoding.
    ///
    /// Useful for hand-crafted or deliberately malformed wire data, e.g. a
    /// fixed header that announces a body which is never sent.
    ///
    /// # Panics
    ///
    /// Panics if the socket write fails.
    pub async fn write_raw(&mut self, bytes: &[u8]) {
        self.try_write_raw(bytes)
            .await
            .expect("fake broker write failed");
    }

    /// Writes raw bytes to the client, returning any socket error instead of
    /// panicking.
    ///
    /// Use this when the client is expected to apply backpressure or close
    /// the connection, so the write failing (or never returning) is part of
    /// what the test observes.
    ///
    /// # Errors
    ///
    /// Returns the underlying I/O error if the socket write fails.
    pub async fn try_write_raw(&mut self, bytes: &[u8]) -> io::Result<()> {
        self.stream.write_all(bytes).await
    }

    /// Reads the next packet and asserts it is a CONNECT.
    ///
    /// # Panics
    ///
    /// Panics if the next packet is not a CONNECT.
    pub async fn expect_connect(&mut self) -> RecvPacket {
        let pkt = self.read_packet().await;
        assert_eq!(
            pkt.packet_type(),
            PacketType::Connect,
            "expected CONNECT as the first packet"
        );
        pkt
    }

    /// Sends a successful CONNACK with the given properties.
    ///
    /// # Panics
    ///
    /// Panics if a property is not valid for CONNACK.
    pub async fn send_connack(&mut self, session_present: bool, props: &[Property<'_>]) {
        let buf = ConnackPropertiesBuf::try_from(props).expect("invalid CONNACK property");
        self.write_packet(&Packet::Connack(Connack {
            session_present,
            reason_code: ReasonCode::Success,
            properties: buf.as_properties(),
        }))
        .await;
    }

    /// Performs the broker side of the connection handshake: reads CONNECT,
    /// replies with a successful CONNACK (`session_present = false`) carrying
    /// the given properties.
    pub async fn handshake(&mut self, props: &[Property<'_>]) {
        self.expect_connect().await;
        self.send_connack(false, props).await;
    }

    /// Acknowledges a `QoS` 1 PUBLISH with a successful PUBACK.
    pub async fn send_puback(&mut self, packet_id: PacketId) {
        self.write_packet(&Packet::Puback(Ack {
            packet_id,
            reason_code: ReasonCode::Success,
            properties: AckProperties::EMPTY,
        }))
        .await;
    }

    /// Acknowledges a SUBSCRIBE with one reason code per requested filter.
    pub async fn send_suback(&mut self, packet_id: PacketId, reason_codes: &[ReasonCode]) {
        self.write_packet(&Packet::Suback(Suback {
            packet_id,
            properties: SubackProperties::EMPTY,
            reason_codes,
        }))
        .await;
    }

    /// Reads packets until a PUBLISH arrives, transparently answering
    /// PINGREQ with PINGRESP. Returns the PUBLISH packet.
    ///
    /// # Panics
    ///
    /// Panics if a packet other than PUBLISH or PINGREQ arrives, or if no
    /// PUBLISH arrives within 10 s.
    pub async fn expect_publish(&mut self) -> RecvPacket {
        self.expect_packet(PacketType::Publish).await
    }

    /// Reads packets until a SUBSCRIBE arrives, transparently answering
    /// PINGREQ with PINGRESP. Returns the SUBSCRIBE packet.
    ///
    /// # Panics
    ///
    /// Panics if a packet other than SUBSCRIBE or PINGREQ arrives, or if no
    /// SUBSCRIBE arrives within 10 s.
    pub async fn expect_subscribe(&mut self) -> RecvPacket {
        self.expect_packet(PacketType::Subscribe).await
    }

    /// Reads packets until one of type `want` arrives, transparently
    /// answering PINGREQ with PINGRESP.
    ///
    /// # Panics
    ///
    /// Panics if a packet other than `want` or PINGREQ arrives, or if no
    /// matching packet arrives within 10 s.
    async fn expect_packet(&mut self, want: PacketType) -> RecvPacket {
        loop {
            let pkt = self.read_packet().await;
            let got = pkt.packet_type();
            if got == want {
                return pkt;
            }
            match got {
                PacketType::Pingreq => self.write_packet(&Packet::Pingresp).await,
                other => panic!("expected {want:?}, got {other:?}"),
            }
        }
    }
}