eventloop.rs

use super::{
    framed::Network, mqttbytes::v5::*, Incoming, MqttOptions, MqttState, Outgoing, Request,
    StateError, Transport,
};
use crate::{eventloop::socket_connect, framed::N};

use flume::{bounded, Receiver, Sender};
use futures::{stream::repeat, FutureExt, SinkExt, Stream, StreamExt};
use tokio::{
    select, task,
    time::{self, error::Elapsed, Interval},
};
use tokio_stream::wrappers::IntervalStream;

use std::{collections::VecDeque, convert::TryInto, future::pending, io, time::Duration};

use super::mqttbytes::v5::ConnectReturnCode;

#[cfg(any(feature = "use-rustls", feature = "use-native-tls"))]
use crate::tls;

#[cfg(unix)]
use {std::path::Path, tokio::net::UnixStream};

#[cfg(feature = "websocket")]
use {
    crate::websockets::{split_url, validate_response_headers, UrlError},
    async_tungstenite::tungstenite::client::IntoClientRequest,
    ws_stream_tungstenite::WsStream,
};

#[cfg(feature = "proxy")]
use crate::proxy::ProxyError;

type StateEventSender = Sender<Event>;
type ConnectionEventSender = Sender<Result<Event, ConnectionError>>;
type RequestReceiver = Receiver<Request>;

const MAX_LOOP: usize = 20;

/// Critical errors during eventloop polling
#[derive(Debug, thiserror::Error)]
pub enum ConnectionError {
    #[error("Mqtt state: {0}")]
    MqttState(#[from] StateError),
    #[error("Timeout")]
    Timeout(#[from] Elapsed),
    #[cfg(feature = "websocket")]
    #[error("Websocket: {0}")]
    Websocket(#[from] async_tungstenite::tungstenite::error::Error),
    #[cfg(feature = "websocket")]
    #[error("Websocket Connect: {0}")]
    WsConnect(#[from] http::Error),
    #[cfg(any(feature = "use-rustls", feature = "use-native-tls"))]
    #[error("TLS: {0}")]
    Tls(#[from] tls::Error),
    #[error("I/O: {0}")]
    Io(#[from] io::Error),
    #[error("Connection refused, return code: `{0:?}`")]
    ConnectionRefused(ConnectReturnCode),
    #[error("Expected ConnAck packet, received: {0:?}")]
    NotConnAck(Box<Packet>),
    #[error("Requests done")]
    RequestsDone,
    #[error("Connection closed")]
    Closed,
    #[cfg(feature = "websocket")]
    #[error("Invalid Url: {0}")]
    InvalidUrl(#[from] UrlError),
    #[cfg(feature = "proxy")]
    #[error("Proxy Connect: {0}")]
    Proxy(#[from] ProxyError),
    #[cfg(feature = "websocket")]
    #[error("Websocket response validation error: ")]
    ResponseValidation(#[from] crate::websockets::ValidationError),
}

/// Eventloop with all the state of a connection
pub struct EventLoop {
    // Receiver handle for state events
    state_event_rx: Receiver<Event>,
    // Receiver handle for connection events
    connection_event_rx: Receiver<Result<Event, ConnectionError>>,
}

/// Events which can be yielded by the event loop
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Event {
    Incoming(Incoming),
    Outgoing(Outgoing),
}

impl EventLoop {
    /// New MQTT `EventLoop`
    pub fn new(options: MqttOptions, request_rx: RequestReceiver) -> EventLoop {
        let (state_event_tx, state_event_rx) = bounded(request_rx.capacity().unwrap());
        let (connection_event_tx, connection_event_rx) = bounded(10);

        // Spawn a task that runs the eventloop
        task::spawn(loop_task(
            options,
            request_rx,
            state_event_tx,
            connection_event_tx,
        ));

        EventLoop {
            state_event_rx,
            connection_event_rx,
        }
    }

    /// Poll for the next event
    pub async fn poll(&self) -> Result<Event, ConnectionError> {
        select! {
            event = self.state_event_rx.recv_async() => event.map_err(|_| ConnectionError::Closed),
            event = self.connection_event_rx.recv_async() => event.map_err(|_| ConnectionError::Closed)?,
        }
    }
}

/// Main eventloop task
async fn loop_task(
    options: MqttOptions,
    request_rx: RequestReceiver,
    state_event_sender: StateEventSender,
    connection_event_sender: ConnectionEventSender,
) {
    let mut options = options;
    let mut request_rx = request_rx;
    let inflight_limit = options.outgoing_inflight_upper_limit.unwrap_or(u16::MAX);
    let manual_acks = options.manual_acks;
    let mut pending_requests = VecDeque::new();
    let mut state = MqttState::new(inflight_limit, manual_acks, state_event_sender);

    'outer: loop {
        // Connect.
        let mut io = match connect(&mut options, &mut state).await {
            Ok(network) => network,
            Err(e) => {
                if connection_event_sender.send_async(Err(e)).await.is_err() {
                    // EventLoop handle is dropped.
                    break 'outer;
                }
                continue 'outer;
            }
        };

        // Setup keepalive ticker
        let mut keepalive = time::interval(options.keep_alive);
        // Discard first tick as it's instant
        keepalive.tick().await;

        // Create pending throttle interval
        let mut pending_interval = if options.pending_throttle().is_zero() {
            // No throttle
            repeat(()).left_stream()
        } else {
            let interval = time::interval(options.pending_throttle());
            IntervalStream::new(interval).map(drop).right_stream()
        };

        // Request and packet processing loop
        'select: loop {
            let result = select(
                &mut state,
                &mut request_rx,
                &mut io,
                &mut pending_requests,
                &mut pending_interval,
                &mut keepalive,
            )
            .await;

            match result {
                Ok(()) => (),
                Err(ConnectionError::RequestsDone) => break 'outer,
                Err(e) => {
                    connection_event_sender.send_async(Err(e)).await.ok();
                    break 'select;
                }
            }
        }

        // Request that are in the channel shall be processed with a throttle delay
        // when the next connection is established. Save them to the pending queue.
        pending_requests.extend(request_rx.drain());
    }
}

async fn select<S>(
    state: &mut MqttState,
    requests: &mut Receiver<Request>,
    mut io: &mut Network,
    pending_requests: &mut VecDeque<Request>,
    pending_interval: &mut S,
    keepalive: &mut Interval,
) -> Result<(), ConnectionError>
where
    S: Stream<Item = ()> + Unpin,
{
    // Loop over request and packet processing. This enables grouped writes of outgoing packets
    // that are fed into the network sink and flushed after this loop.
    for _ in 0..MAX_LOOP {
        let is_inflight_full = state.inflight >= state.max_outgoing_inflight;
        let collision = state.collision.is_some();

        // Check if we can process a request from the pending queue
        let no_pending_requests = pending_requests.is_empty();

        // Create a future that yields the next request from the pending queue
        // or resolves never.
        let pending_fut = if no_pending_requests || is_inflight_full || collision {
            pending::<Request>().left_future()
        } else {
            pending_interval
                .next()
                .map(|_| pending_requests.pop_front().unwrap())
                .right_future()
        };

        select! {
            // Process pending requests
            request = pending_fut => state.handle_outgoing_packet(&mut io, request).await?,
            // Process request from the request channel
            request = &mut requests.recv_async(), if no_pending_requests && !is_inflight_full && !collision => match request {
                Ok(request) => state.handle_outgoing_packet(&mut io, request).await?,
                Err(_) => return Err(ConnectionError::RequestsDone),
            },
            // Process network packet
            packet = io.next() => match packet {
                Some(packet) => state.handle_incoming_packet(&mut io, packet?).await?,
                None => return Err(ConnectionError::Closed),
            },
            // We generate pings irrespective of network activity. This keeps the ping logic
            // simple. We can change this behavior in future if necessary (to prevent extra pings)
            _ = keepalive.tick() => {
                state.handle_outgoing_packet(&mut io, Request::PingReq).await?;
                // Flush directly after pingreq
                io.flush().await?;
            }
        }

        // Check if this bulk operation may have finished early: no requests. no network packet.
        if pending_requests.is_empty() && requests.is_empty() && !io.is_data_pending() {
            break;
        }
    }

    io.flush().await?;

    Ok(())
}

async fn connect(
    options: &mut MqttOptions,
    state: &mut MqttState,
) -> Result<Network, ConnectionError> {
    let timeout = Duration::from_secs(options.connection_timeout());
    let connect = async move {
        let mut network = network_connect(options).await?;
        let packet = mqtt_connect(options, &mut network).await?;
        Ok::<_, ConnectionError>((network, packet))
    };

    let (mut network, connack) = time::timeout(timeout, connect).await??;

    state.handle_incoming_packet(&mut network, connack).await?;
    network.flush().await?;

    Ok(network)
}

async fn network_connect(options: &MqttOptions) -> Result<Network, ConnectionError> {
    let mut max_incoming_pkt_size = Some(options.default_max_incoming_size);

    // Override default value if max_packet_size is set on `connect_properties`
    if let Some(connect_props) = &options.connect_properties {
        if let Some(max_size) = connect_props.max_packet_size {
            let max_size = max_size.try_into().map_err(StateError::Coversion)?;
            max_incoming_pkt_size = Some(max_size);
        }
    }

    // Process Unix files early, as proxy is not supported for them.
    #[cfg(unix)]
    if matches!(options.transport(), Transport::Unix) {
        let file = options.broker_addr.as_str();
        let socket = UnixStream::connect(Path::new(file)).await?;
        let network = Network::new(socket, max_incoming_pkt_size);
        return Ok(network);
    }

    // For websockets domain and port are taken directly from `broker_addr` (which is a url).
    let (domain, port) = match options.transport() {
        #[cfg(feature = "websocket")]
        Transport::Ws => split_url(&options.broker_addr)?,
        #[cfg(all(feature = "use-rustls", feature = "websocket"))]
        Transport::Wss(_) => split_url(&options.broker_addr)?,
        _ => options.broker_address(),
    };

    let tcp_stream: Box<dyn N> = {
        #[cfg(feature = "proxy")]
        match options.proxy() {
            Some(proxy) => {
                proxy
                    .connect(&domain, port, options.network_options())
                    .await?
            }
            None => {
                let addr = format!("{domain}:{port}");
                let tcp = socket_connect(addr, options.network_options()).await?;
                Box::new(tcp)
            }
        }
        #[cfg(not(feature = "proxy"))]
        {
            let addr = format!("{domain}:{port}");
            let tcp = socket_connect(addr, options.network_options()).await?;
            Box::new(tcp)
        }
    };

    let network = match options.transport() {
        Transport::Tcp => Network::new(tcp_stream, max_incoming_pkt_size),
        #[cfg(any(feature = "use-native-tls", feature = "use-rustls"))]
        Transport::Tls(tls_config) => {
            let socket =
                tls::tls_connect(&options.broker_addr, options.port, &tls_config, tcp_stream)
                    .await?;
            Network::new(socket, max_incoming_pkt_size)
        }
        #[cfg(unix)]
        Transport::Unix => unreachable!(),
        #[cfg(feature = "websocket")]
        Transport::Ws => {
            let mut request = options.broker_addr.as_str().into_client_request()?;
            request
                .headers_mut()
                .insert("Sec-WebSocket-Protocol", "mqtt".parse().unwrap());

            if let Some(request_modifier) = options.request_modifier() {
                request = request_modifier(request).await;
            }

            let (socket, response) =
                async_tungstenite::tokio::client_async(request, tcp_stream).await?;
            validate_response_headers(response)?;

            Network::new(WsStream::new(socket), max_incoming_pkt_size)
        }
        #[cfg(all(feature = "use-rustls", feature = "websocket"))]
        Transport::Wss(tls_config) => {
            let mut request = options.broker_addr.as_str().into_client_request()?;
            request
                .headers_mut()
                .insert("Sec-WebSocket-Protocol", "mqtt".parse().unwrap());

            if let Some(request_modifier) = options.request_modifier() {
                request = request_modifier(request).await;
            }

            let connector = tls::rustls_connector(&tls_config).await?;

            let (socket, response) = async_tungstenite::tokio::client_async_tls_with_connector(
                request,
                tcp_stream,
                Some(connector),
            )
            .await?;
            validate_response_headers(response)?;

            Network::new(WsStream::new(socket), max_incoming_pkt_size)
        }
    };

    Ok(network)
}

async fn mqtt_connect(
    options: &mut MqttOptions,
    network: &mut Network,
) -> Result<Incoming, ConnectionError> {
    let keep_alive = options.keep_alive().as_secs() as u16;
    let clean_start = options.clean_start();
    let client_id = options.client_id();
    let properties = options.connect_properties();

    let connect = Connect {
        keep_alive,
        client_id,
        clean_start,
        properties,
    };

    // Send mqtt connect packet
    network.connect(connect, options).await?;

    // Validate connack
    match network.next().await {
        Some(Ok(Incoming::ConnAck(connack))) if connack.code == ConnectReturnCode::Success => {
            // Override local keep_alive value if set by server.
            if let Some(props) = &connack.properties {
                if let Some(keep_alive) = props.server_keep_alive {
                    options.keep_alive = Duration::from_secs(keep_alive as u64);
                }
            }
            Ok(Packet::ConnAck(connack))
        }
        Some(Ok(Incoming::ConnAck(connack))) => {
            Err(ConnectionError::ConnectionRefused(connack.code))
        }
        Some(Ok(packet)) => Err(ConnectionError::NotConnAck(Box::new(packet))),
        _ => unimplemented!(),
    }
}