wireframe 0.3.0

Simplify building servers and clients for custom binary protocols.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! Preamble exchange logic for wireframe client connections.
//!
//! This module handles writing the client preamble to the server and invoking
//! success/failure callbacks during the connection handshake phase.

use std::{io, sync::Arc, time::Duration};

use bincode::Encode;
use futures::future::BoxFuture;
use log::error;
use tokio::{net::TcpStream, time::timeout};

use super::{ClientError, ClientPreambleFailureHandler, ClientPreambleSuccessHandler};
use crate::preamble::write_preamble;

/// Holds optional preamble configuration for the client builder.
pub(crate) struct PreambleConfig<P> {
    pub(crate) preamble: P,
    pub(crate) on_success: Option<ClientPreambleSuccessHandler<P>>,
    pub(crate) on_failure: Option<ClientPreambleFailureHandler>,
    pub(crate) timeout: Option<Duration>,
}

impl<P> Clone for PreambleConfig<P>
where
    P: Clone,
{
    fn clone(&self) -> Self {
        Self {
            preamble: self.preamble.clone(),
            on_success: self.on_success.clone(),
            on_failure: self.on_failure.clone(),
            timeout: self.timeout,
        }
    }
}

impl<P> PreambleConfig<P>
where
    P: Encode + Send + Sync + 'static,
{
    /// Create a new preamble configuration with just the preamble value.
    pub(crate) fn new(preamble: P) -> Self {
        Self {
            preamble,
            on_success: None,
            on_failure: None,
            timeout: None,
        }
    }

    /// Set the timeout for the preamble exchange.
    pub(crate) fn set_timeout(&mut self, timeout: Duration) { self.timeout = Some(timeout); }

    /// Set the success handler.
    pub(crate) fn set_success_handler<H>(&mut self, handler: H)
    where
        H: for<'a> Fn(&'a P, &'a mut TcpStream) -> BoxFuture<'a, io::Result<Vec<u8>>>
            + Send
            + Sync
            + 'static,
    {
        self.on_success = Some(Arc::new(handler));
    }

    /// Set the failure handler.
    pub(crate) fn set_failure_handler<H>(&mut self, handler: H)
    where
        H: for<'a> Fn(&'a ClientError, &'a mut TcpStream) -> BoxFuture<'a, io::Result<()>>
            + Send
            + Sync
            + 'static,
    {
        self.on_failure = Some(Arc::new(handler));
    }
}

/// Perform the preamble exchange: write preamble, invoke success callback.
pub(crate) async fn perform_preamble_exchange<P>(
    stream: &mut TcpStream,
    config: PreambleConfig<P>,
) -> Result<Vec<u8>, ClientError>
where
    P: Encode + Send + Sync + 'static,
{
    let PreambleConfig {
        preamble,
        on_success,
        on_failure,
        timeout: preamble_timeout,
    } = config;

    let result = run_preamble_exchange(stream, &preamble, on_success, preamble_timeout).await;

    // On failure, invoke the failure callback if registered.
    if let Err(ref err) = result {
        invoke_failure_handler(stream, err, on_failure.as_ref()).await;
    }

    result
}

/// Execute the preamble write and success callback, with optional timeout.
async fn run_preamble_exchange<P>(
    stream: &mut TcpStream,
    preamble: &P,
    on_success: Option<ClientPreambleSuccessHandler<P>>,
    preamble_timeout: Option<Duration>,
) -> Result<Vec<u8>, ClientError>
where
    P: Encode + Send + Sync + 'static,
{
    let exchange = async {
        // Write the preamble.
        write_preamble(stream, preamble)
            .await
            .map_err(ClientError::PreambleEncode)?;

        // Invoke success callback if registered, collecting leftover bytes.
        match on_success.as_ref() {
            Some(handler) => handler(preamble, stream)
                .await
                .map_err(ClientError::PreambleRead),
            None => Ok(Vec::new()),
        }
    };

    // Apply timeout if configured.
    match preamble_timeout {
        Some(limit) => timeout(limit, exchange)
            .await
            .unwrap_or(Err(ClientError::PreambleTimeout)),
        None => exchange.await,
    }
}

/// Invoke the failure handler if one is registered.
async fn invoke_failure_handler(
    stream: &mut TcpStream,
    err: &ClientError,
    on_failure: Option<&ClientPreambleFailureHandler>,
) {
    if let Some(handler) = on_failure
        && let Err(e) = handler(err, stream).await
    {
        error!("preamble failure handler error: {e}");
    }
}