ecksport_net/

handshake.rs

//! Subroutines for performing the connection setup, including managing any
//! challenges.

use std::io;
use std::time;

use thiserror::Error;
use tracing::*;

use ecksport_core::frame::{FrameBody, FrameType, GoodbyeData, ServerDumpData};
use ecksport_core::peer::PeerData;
use ecksport_core::state_mach::{self, Action, ClientMeta, ReadyData, ServerMeta, ServerState};
use ecksport_core::traits::AuthConfig;
use ecksport_core::traits::{AsyncRecvFrame, AsyncSendFrame};
use ecksport_core::{errors::*, topic};

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

impl HandshakeOptions {
    pub fn new(timeout: time::Duration) -> Self {
        Self { timeout }
    }
}

#[derive(Debug, Error)]
pub enum HandshakeError {
    #[error("timed out")]
    Timeout,

    #[error("we closed connection: {0:?}")]
    Abort(GoodbyeData),

    #[error("connection closed on us: {0:?}")]
    Closed(GoodbyeData),

    #[error("we're happy and we don't know why")]
    ClosedUnclear,

    #[error("unexpected frame type {0:?}")]
    UnexpectedFrameType(FrameType),

    #[error("conn: {0}")]
    Conn(#[from] ConnError),

    #[error("io: {0}")]
    Io(#[from] io::Error),
}

/// Contains data about how a handshake completed, which resulted in a ready connection.
#[derive(Clone, Debug)]
pub struct HandshakeStatus {
    ready_data: ReadyData,
    peer_data: PeerData,
}

impl HandshakeStatus {
    fn new(ready_data: ReadyData, peer_data: PeerData) -> Self {
        Self {
            ready_data,
            peer_data,
        }
    }

    pub fn ready(&self) -> &ReadyData {
        &self.ready_data
    }

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

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

pub async fn do_client_handshake_async<
    T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin,
    A: AuthConfig,
>(
    conn: &mut T,
    protocol: topic::Topic,
    meta: &ClientMeta,
    opts: &HandshakeOptions,
    auth: A,
    mut peer: PeerData,
) -> Result<HandshakeStatus, HandshakeError> {
    let (mut cstate, mut actions) = state_mach::exec_client_init(protocol, auth, meta);

    loop {
        // Perform the last action, then wait for an input.
        match perform_actions(&actions, conn, &mut peer).await? {
            ActionStatus::FinishReady(rd) => {
                let hs = HandshakeStatus::new(rd, peer);
                return Ok(hs);
            }
            ActionStatus::FinishExit => return Err(HandshakeError::ClosedUnclear),
            _ => {}
        }

        // Recv message or timeout.
        let m = match tokio::time::timeout(opts.timeout, conn.recv_frame_async()).await {
            Ok(v) => v?,
            Err(_) => return Err(HandshakeError::Timeout),
        };

        // Update the state and decide what to do with it.
        (cstate, actions) = state_mach::exec_client_inp(&cstate, &m, meta);
    }
}

pub async fn do_server_handshake_async<
    T: AsyncSendFrame + AsyncRecvFrame + Sync + Send + Unpin,
    A: AuthConfig,
>(
    conn: &mut T,
    meta: &ServerMeta,
    opts: &HandshakeOptions,
    auth: A,
    mut peer: PeerData,
) -> Result<Option<HandshakeStatus>, HandshakeError> {
    let mut sstate = ServerState::init(auth);

    loop {
        // Recv message or timeout.
        let m = match tokio::time::timeout(opts.timeout, conn.recv_frame_async()).await {
            Ok(v) => v?,
            Err(_) => return Err(HandshakeError::Timeout),
        };

        let (new_state, actions) = state_mach::exec_server_inp(&sstate, &m, meta);
        sstate = new_state;

        // Perform the action we've decided to do.
        match perform_actions(&actions, conn, &mut peer).await? {
            ActionStatus::FinishReady(rd) => {
                let hs = HandshakeStatus::new(rd, peer);
                return Ok(Some(hs));
            }
            ActionStatus::FinishExit => return Ok(None),
            _ => {}
        }
    }
}

/// Describes how we should proceed from an action being performed.  This is
/// interpreted differently depending on if we're the client or the server.
pub enum ActionStatus {
    Continue,
    FinishReady(ReadyData),
    FinishExit,
}

async fn perform_actions<T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin>(
    actions: impl IntoIterator<Item = &Action>,
    iow: &mut T,
    peer: &mut PeerData,
) -> Result<ActionStatus, HandshakeError> {
    let mut status = ActionStatus::Continue;

    for a in actions {
        let res = perform_action(a, iow, peer).await?;
        if !matches!(res, ActionStatus::Continue) {
            status = res;
        }
    }

    Ok(status)
}

async fn perform_action<T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin>(
    action: &Action,
    iow: &mut T,
    peer: &mut PeerData,
) -> Result<ActionStatus, HandshakeError> {
    match action {
        Action::SendFrame(frame) => {
            iow.send_frame_async(&frame).await?;
            Ok(ActionStatus::Continue)
        }

        Action::Ready(rd) => {
            trace!("readying connection");
            return Ok(ActionStatus::FinishReady(rd.clone()));
        }

        Action::Abort(gd) => {
            let fb = FrameBody::Goodbye(gd.clone());
            iow.send_frame_async(&fb).await?;
            Err(HandshakeError::Abort(gd.clone()))
        }

        Action::Exit(gd) => Err(HandshakeError::Closed(gd.clone())),

        Action::SetRemoteAgent(agent) => {
            peer.set_agent(agent.clone());
            Ok(ActionStatus::Continue)
        }

        Action::SetRemoteIdent(ident) => {
            peer.set_identity(ident.clone());
            Ok(ActionStatus::Continue)
        }

        Action::HappyClose => Ok(ActionStatus::FinishExit),
    }
}

pub async fn do_query_async<T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin>(
    conn: &mut T,
    meta: &ClientMeta,
    opts: &HandshakeOptions,
) -> Result<ServerDumpData, HandshakeError> {
    // First send a client query frame.
    let frame = meta.create_query_frame();
    conn.send_frame_async(&frame).await?;

    // Recv message or timeout.
    let m = match tokio::time::timeout(opts.timeout, conn.recv_frame_async()).await {
        Ok(v) => v?,
        Err(_) => return Err(HandshakeError::Timeout),
    };

    // If it's what we expected then we can return it, otherwise error.
    match m {
        FrameBody::ServerDump(sdump) => Ok(sdump),
        _ => Err(HandshakeError::UnexpectedFrameType(m.ty())),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::net;

    use ecksport_core::{stream_framing, topic};
    use tokio::task::JoinHandle;

    const DUMMY_TOPIC: topic::Topic = topic::Topic::from_const_str("dummytop");

    #[tokio::test]
    async fn test_handshake() {
        let cmeta = ClientMeta::new("/test/client/".to_owned());
        let smeta = ServerMeta::new("/test/server/".to_owned(), vec![DUMMY_TOPIC]);
        let opts = HandshakeOptions {
            timeout: time::Duration::from_millis(1000),
        };
        let opts2 = opts.clone();

        // Random port that we picked.
        let addr = "127.0.0.1:31827"
            .parse::<net::SocketAddr>()
            .expect("test: parse socketaddr");

        let sock = tokio::net::TcpListener::bind(addr)
            .await
            .expect("test: tcp listen");

        let shj: JoinHandle<Result<(), HandshakeError>> = tokio::spawn(async move {
            let (conn, _) = sock.accept().await?;
            let mut framer = stream_framing::StreamFramer::new(conn);
            let pd = PeerData::default();
            do_server_handshake_async(&mut framer, &smeta, &opts, (), pd).await?;
            Ok(())
        });

        let chj: JoinHandle<Result<(), HandshakeError>> = tokio::spawn(async move {
            let conn = tokio::net::TcpStream::connect(addr).await?;
            let mut framer = stream_framing::StreamFramer::new(conn);
            let pd = PeerData::default();
            do_client_handshake_async(&mut framer, DUMMY_TOPIC, &cmeta, &opts2, (), pd).await?;
            Ok(())
        });

        shj.await
            .expect("test: server wait")
            .expect("test: server connect");

        chj.await
            .expect("test: client wait")
            .expect("test: client connect");
    }

    #[tokio::test]
    async fn test_query() {
        let cmeta = ClientMeta::new("/test/client/".to_owned());
        let smeta = ServerMeta::new("/test/server/".to_owned(), vec![DUMMY_TOPIC]);
        let opts = HandshakeOptions {
            timeout: time::Duration::from_millis(1000),
        };
        let opts2 = opts.clone();

        // Random port that we picked.
        let addr = "127.0.0.1:31828"
            .parse::<net::SocketAddr>()
            .expect("test: parse socketaddr");

        let sock = tokio::net::TcpListener::bind(addr)
            .await
            .expect("test: tcp listen");

        let shj: JoinHandle<Result<(), HandshakeError>> = tokio::spawn(async move {
            let (conn, _) = sock.accept().await?;
            let mut framer = stream_framing::StreamFramer::new(conn);
            let pd = PeerData::default();
            do_server_handshake_async(&mut framer, &smeta, &opts, (), pd).await?;
            Ok(())
        });

        let chj: JoinHandle<Result<ServerDumpData, HandshakeError>> = tokio::spawn(async move {
            let conn = tokio::net::TcpStream::connect(addr).await?;
            let mut framer = stream_framing::StreamFramer::new(conn);

            let sdump = do_query_async(&mut framer, &cmeta, &opts2).await?;
            Ok(sdump)
        });

        let _ = shj.await.expect("test: server wait");

        let sdump = chj
            .await
            .expect("test: client wait")
            .expect("test: client connect");

        assert_eq!(sdump.agent(), "/test/server/");
        assert_eq!(sdump.protocols(), &[DUMMY_TOPIC]);
    }
}