ecksport_net/

connection.rs

//! Connection bookkeeping to track and emit user-facing events.

use std::collections::*;
use std::time;

use ecksport_core::peer::PeerData;

use ecksport_core::frame::{CloseData, FrameBody, MsgFlags, NotificationData, OpenData, PushData};
use ecksport_core::state_mach::{ClientMeta, ServerMeta};
use ecksport_core::topic;
use ecksport_core::traits::{AsyncRecvFrame, AsyncSendFrame, AuthConfig};

use crate::channel_state::{self, Creator};
use crate::errors::Error;
use crate::event::{InbEvent, PushFlags};
use crate::handshake::{self, do_server_handshake_async};

/// Low-level connection type that handles basic channel idx bookkeeping and
/// inbound event buffering.
pub struct Connection<T> {
    inner: T,

    /// Protocol chosen by initiator.
    protocol: topic::Topic,

    /// The side that created the channel.
    initiator: Creator,

    /// Identity data about the peer figured out during handshake.
    peer: PeerData,

    /// Bookkeeping for channel state.
    chan_tbl: channel_state::ChannelTable,

    /// Buffer for events that have been produced in response to frames but not
    /// emitted yet.
    event_queue: VecDeque<InbEvent>,
}

impl<T> Connection<T> {
    /// Wraps an existing connection that's already been handshaked.
    fn new(inner: T, initiator: Creator, protocol: topic::Topic, peer: PeerData) -> Self {
        Self {
            inner,
            protocol,
            initiator,
            peer,
            chan_tbl: channel_state::ChannelTable::new(initiator),
            event_queue: VecDeque::new(),
        }
    }

    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Returns a mutable reference to the inner recv instance.  Note that using
    /// this to recv messages can break the bookkeeping, so use caution.
    pub fn inner_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    pub fn into_inner(self) -> T {
        self.inner
    }

    pub fn protocol(&self) -> topic::Topic {
        self.protocol
    }

    pub fn initiator(&self) -> Creator {
        self.initiator
    }

    pub fn peer_data(&self) -> &PeerData {
        &self.peer
    }

    /// Returns the number of open channels in the internal channel table.  This
    /// includes channels opened in response to pending events that are buffered
    /// and need to be processed.
    pub fn num_open_channels(&self) -> usize {
        self.chan_tbl.num_open_channels()
    }

    /// Returns if there's events that have been produced but not consumed yet.
    pub fn has_pending_events(&self) -> bool {
        !self.event_queue.is_empty()
    }

    fn handle_frame(&mut self, frame: FrameBody) -> Result<(), Error> {
        match frame {
            FrameBody::OpenChan(open) => {
                // Set up new channel bookkeeping.
                let close = open.close();
                let flags = PushFlags::from(open.flags());
                let id = self.chan_tbl.init_remote_chan(open.topic(), close);

                // Insert appropriate events.
                self.event_queue.push_back(InbEvent::NewChannel(
                    id,
                    open.topic(),
                    flags,
                    open.into_payload(),
                ));
                if close {
                    self.event_queue.push_back(InbEvent::CloseChannel(id, true))
                }

                Ok(())
            }

            FrameBody::PushChan(push) => {
                let id = push.chan_id();
                self.chan_tbl.check_recv_on_chan(id)?;
                let flags = PushFlags::from(push.flags());
                let close = push.close();

                self.event_queue
                    .push_back(InbEvent::PushChannel(id, flags, push.into_payload()));

                if close {
                    let removed = self
                        .chan_tbl
                        .mark_chan_remote_closed(id)
                        .expect("connection: close remote");

                    self.event_queue
                        .push_back(InbEvent::CloseChannel(id, !removed));
                }

                Ok(())
            }

            FrameBody::CloseChan(close) => {
                self.chan_tbl.check_recv_on_chan(close.chan_id())?;

                let removed = self
                    .chan_tbl
                    .mark_chan_remote_closed(close.chan_id())
                    .expect("connection: close remote");

                self.event_queue
                    .push_back(InbEvent::CloseChannel(close.chan_id(), !removed));
                Ok(())
            }

            FrameBody::Notification(notif) => {
                let topic = notif.topic();
                self.event_queue
                    .push_back(InbEvent::Notification(topic, notif.into_payload()));
                Ok(())
            }

            _ => Err(Error::UnexpectedFrame(frame.ty())),
        }
    }
}

impl<T: AsyncRecvFrame> Connection<T> {
    async fn recv_frame(&mut self) -> Result<(), Error> {
        let frame = self.inner.recv_frame_async().await?;
        self.handle_frame(frame)?;
        Ok(())
    }

    /// Takes the next event out of the buffer or, if empty, waits for a new
    /// event to be produced.
    pub async fn next_event(&mut self) -> Result<Option<InbEvent>, Error> {
        if !self.event_queue.is_empty() {
            return Ok(Some(self.event_queue.pop_front().unwrap()));
        }

        self.recv_frame().await?;
        Ok(self.event_queue.pop_front())
    }
}

impl<T: AsyncSendFrame> Connection<T> {
    /// Opens a channel on a topic with a channel, optionally closing our end of
    /// it immediately.
    pub async fn open_channel(
        &mut self,
        topic: topic::Topic,
        payload: Vec<u8>,
        flags: MsgFlags,
    ) -> Result<u32, Error> {
        let im_close = flags.close;
        let open_data = OpenData::new(topic, flags, payload);
        let frame = FrameBody::OpenChan(open_data);
        self.inner.send_frame_async(&frame).await?;
        let id = self.chan_tbl.init_local_chan(topic, !im_close);
        Ok(id)
    }

    /// Sends a message on a channel that we haven't closed ourselves yet.
    /// Returns if the channel is kept alive after this.
    pub async fn send_message(
        &mut self,
        chan_id: u32,
        payload: Vec<u8>,
        flags: MsgFlags,
    ) -> Result<bool, Error> {
        self.chan_tbl.check_send_on_chan(chan_id)?;

        let push_data = PushData::new(chan_id, flags, payload);
        let frame = FrameBody::PushChan(push_data);
        self.inner.send_frame_async(&frame).await?;

        if flags.close {
            let removed = self
                .chan_tbl
                .mark_chan_local_closed(chan_id)
                .expect("connection: close local");
            Ok(!removed)
        } else {
            Ok(true)
        }
    }

    /// Close a channel without sending a message.  Returns if the channel is
    /// kept alive after this.
    pub async fn close_channel(&mut self, chan_id: u32) -> Result<bool, Error> {
        self.chan_tbl.check_send_on_chan(chan_id)?;

        let close_data = CloseData::new(chan_id);
        let frame = FrameBody::CloseChan(close_data);
        self.inner.send_frame_async(&frame).await?;

        let removed = self
            .chan_tbl
            .mark_chan_local_closed(chan_id)
            .expect("connection: close local");

        Ok(!removed)
    }

    /// Sends a notification on a topic.
    pub async fn send_notification(
        &mut self,
        topic: topic::Topic,
        message: Vec<u8>,
    ) -> Result<(), Error> {
        let notif_data = NotificationData::new(topic, message);
        let frame = FrameBody::Notification(notif_data);
        self.inner.send_frame_async(&frame).await?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct ConnectOptions {
    pub timeout: time::Duration,
    pub client_meta: ClientMeta,
}

impl Default for ConnectOptions {
    fn default() -> Self {
        Self {
            timeout: time::Duration::from_millis(15000),
            client_meta: ClientMeta::new("/ecksport/alpha/".to_owned()),
        }
    }
}

/// Takes a freshly constructed connection, available as frames, and performs
/// the client side of a normal handshake to select the specified protocol,
/// returning a new low-level connection if the handshake completes
/// successfully.
pub async fn perform_handshake_async<
    T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static,
    A: AuthConfig,
>(
    mut stream: T,
    protocol: topic::Topic,
    opts: ConnectOptions,
    auth: A,
    peer: PeerData,
) -> Result<Connection<T>, Error> {
    let hs_opts = handshake::HandshakeOptions::new(opts.timeout);

    // Perform the actual client handshake.
    let hs = handshake::do_client_handshake_async(
        &mut stream,
        protocol,
        &opts.client_meta,
        &hs_opts,
        auth,
        peer,
    )
    .await?;
    assert_eq!(hs.ready().protocol(), protocol);

    let peer = hs.into_peer();
    Ok(Connection::new(stream, Creator::Local, protocol, peer))
}

#[derive(Clone, Debug)]
pub struct AcceptOptions {
    pub timeout: time::Duration,
    pub server_meta: ServerMeta,
}

impl Default for AcceptOptions {
    fn default() -> Self {
        Self {
            timeout: time::Duration::from_millis(15000),
            server_meta: ServerMeta::new("/ecksport/alpha/".to_owned(), Vec::new()),
        }
    }
}

/// Takes a freshly accepted connection and performs the server side of a
/// handshake, returning a new low-level connection if the handshake completes
/// successfully.
// TODO figure out how to reuse the options on each accept
pub async fn accept_connection_async<
    T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static,
    A: AuthConfig,
>(
    mut stream: T,
    opts: AcceptOptions,
    auth: A,
    peer: PeerData,
) -> Result<Option<Connection<T>>, Error> {
    let hs_opts = handshake::HandshakeOptions::new(opts.timeout);

    // The server side is simpler.
    let Some(hs) =
        do_server_handshake_async(&mut stream, &opts.server_meta, &hs_opts, auth, peer).await?
    else {
        return Ok(None);
    };

    let proto = hs.ready().protocol();
    let peer = hs.into_peer();

    let conn = Connection::new(stream, Creator::Remote, proto, peer);
    Ok(Some(conn))
}

#[cfg(test)]
mod tests {
    use core::net;

    use ecksport_core::{stream_framing, topic};

    use crate::builder::ClientBuilder;

    use super::*;

    /// Simple connection over a TCP stream.
    pub type TokioTcpConnection = Connection<StreamFramer<tokio::net::TcpStream>>;

    /// Connects to a TCP socket with a Tokio channel, requests to use a particular
    /// protocol, and returns a low-level connection to be used.
    async fn connect_tcp_tokio<A: AuthConfig>(
        socket_addr: SocketAddr,
        protocol: topic::Topic,
        opts: ConnectOptions,
        auth: A,
    ) -> Result<TokioTcpConnection, Error> {
        // Do the initial connection.
        let socket_connect_fut = tokio::net::TcpStream::connect(socket_addr);
        let sock = match timeout(opts.timeout, socket_connect_fut).await {
            Ok(res) => res?,
            Err(_) => return Err(Error::ConnectionTimeout),
        };

        let peer = PeerData::new_loc(Location::Ip(socket_addr));
        let framer = StreamFramer::new(sock);
        Ok(perform_handshake_async(framer, protocol, opts, auth, peer).await?)
    }

    /// Basic connection test that opens and closes a few channels and does some
    /// basic checks that they make sense.
    #[tokio::test]
    async fn test_connect_accept() {
        let socket_addr = "127.0.0.1:5436"
            .parse::<net::SocketAddr>()
            .expect("test: parse addr");

        let lis = tokio::net::TcpListener::bind(socket_addr)
            .await
            .expect("test: bind");

        let proto = topic::Topic::from_const_str("TESTTEST");
        let topic = topic::Topic::from_const_str("FOOOBARR");
        let topic2 = topic::Topic::from_const_str("BAZZQUUX");
        let mut acc_opts = AcceptOptions::default();
        acc_opts.server_meta.add_protocol(proto);
        let conn_opts = ConnectOptions::default();

        let lj = tokio::spawn(async move {
            let (sock, _sa) = lis.accept().await.expect("test: accept");
            let framer = stream_framing::StreamFramer::new(sock);
            let pd = PeerData::default();

            let mut conn = accept_connection_async(framer, acc_opts, (), pd)
                .await
                .expect("test: server handshake")
                .expect("test: create server connection");

            // Should be an open with data.
            let ev = conn
                .next_event()
                .await
                .expect("test: accept event")
                .expect("test: read event");
            eprintln!("got event: {ev:?}");

            assert_eq!(conn.num_open_channels(), 1);

            // Open a new channel.
            conn.open_channel(topic2, vec![5, 6, 7, 8], MsgFlags::none())
                .await
                .expect("test: open channel");

            assert_eq!(conn.num_open_channels(), 2);

            // Should be a close.
            let ev = conn
                .next_event()
                .await
                .expect("test: recv frame")
                .expect("test: recv event");
            eprintln!("got event: {ev:?}");

            assert_eq!(conn.num_open_channels(), 2);

            conn.close_channel(1).await.expect("test: close channel");

            eprintln!("closing channel 0 on the server side");
            conn.close_channel(0)
                .await
                .expect("test: close client chan");
        });

        let cj = tokio::spawn(async move {
            let mut conn = connect_tcp_tokio(socket_addr, proto, conn_opts, ())
                .await
                .expect("test: connect and handshake");

            eprintln!("opening channel, will im_close");
            let ch_idx = conn
                .open_channel(topic, vec![1, 2, 3, 4], MsgFlags::close())
                .await
                .expect("test: open channel");
            assert_eq!(ch_idx, 0);

            assert_eq!(conn.num_open_channels(), 1);

            let ev = conn
                .next_event()
                .await
                .expect("test: recv frame")
                .expect("test: recv event");
            assert_eq!(ev.chan_id(), Some(1));

            assert_eq!(conn.num_open_channels(), 2);
        });

        lj.await.expect("test: server side");
        cj.await.expect("test: client side");

        // test
    }
}