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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
use crate::{Listener, PortRange, TcpTransport};
use async_trait::async_trait;
use std::{fmt, io, net::IpAddr};
use tokio::net::TcpListener as TokioTcpListener;

/// Represents a [`Listener`] for incoming connections over TCP
pub struct TcpListener {
    addr: IpAddr,
    port: u16,
    inner: TokioTcpListener,
}

impl TcpListener {
    /// Creates a new listener by binding to the specified IP address and port
    /// in the given port range
    pub async fn bind(addr: IpAddr, port: impl Into<PortRange>) -> io::Result<Self> {
        let listener =
            TokioTcpListener::bind(port.into().make_socket_addrs(addr).as_slice()).await?;

        // Get the port that we bound to
        let port = listener.local_addr()?.port();

        Ok(Self {
            addr,
            port,
            inner: listener,
        })
    }

    /// Returns the IP address that the listener is bound to
    pub fn ip_addr(&self) -> IpAddr {
        self.addr
    }

    /// Returns the port that the listener is bound to
    pub fn port(&self) -> u16 {
        self.port
    }
}

impl fmt::Debug for TcpListener {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TcpListener")
            .field("addr", &self.addr)
            .field("port", &self.port)
            .finish()
    }
}

#[async_trait]
impl Listener for TcpListener {
    type Output = TcpTransport;

    async fn accept(&mut self) -> io::Result<Self::Output> {
        let (stream, peer_addr) = TokioTcpListener::accept(&self.inner).await?;
        Ok(TcpTransport {
            addr: peer_addr.ip(),
            port: peer_addr.port(),
            inner: stream,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{Ipv6Addr, SocketAddr};
    use tokio::{
        io::{AsyncReadExt, AsyncWriteExt},
        sync::oneshot,
        task::JoinHandle,
    };

    #[tokio::test]
    async fn should_fail_to_bind_if_port_already_bound() {
        let addr = IpAddr::V6(Ipv6Addr::LOCALHOST);
        let port = 0; // Ephemeral port

        // Listen at some port
        let listener = TcpListener::bind(addr, port)
            .await
            .expect("Unexpectedly failed to bind first time");

        // Get the actual port we bound to
        let port = listener.port();

        // Now this should fail as we're already bound to the address and port
        TcpListener::bind(addr, port).await.expect_err(&format!(
            "Unexpectedly succeeded in binding a second time to {}:{}",
            addr, port,
        ));
    }

    #[tokio::test]
    async fn should_be_able_to_receive_connections_and_send_and_receive_data_with_them() {
        let (tx, rx) = oneshot::channel();

        // Spawn a task that will wait for two connections and then
        // return the success or failure
        let task: JoinHandle<io::Result<()>> = tokio::spawn(async move {
            let addr = IpAddr::V6(Ipv6Addr::LOCALHOST);
            let port = 0; // Ephemeral port

            // Listen at the address and port
            let mut listener = TcpListener::bind(addr, port).await?;

            // Send the name back to our main test thread
            tx.send(SocketAddr::from((addr, listener.port())))
                .map_err(|x| io::Error::new(io::ErrorKind::Other, x.to_string()))?;

            // Get first connection
            let mut conn_1 = listener.accept().await?;

            // Send some data to the first connection (12 bytes)
            conn_1.write_all(b"hello conn 1").await?;

            // Get some data from the first connection (14 bytes)
            let mut buf: [u8; 14] = [0; 14];
            let _ = conn_1.read_exact(&mut buf).await?;
            assert_eq!(&buf, b"hello server 1");

            // Get second connection
            let mut conn_2 = listener.accept().await?;

            // Send some data on to second connection (12 bytes)
            conn_2.write_all(b"hello conn 2").await?;

            // Get some data from the second connection (14 bytes)
            let mut buf: [u8; 14] = [0; 14];
            let _ = conn_2.read_exact(&mut buf).await?;
            assert_eq!(&buf, b"hello server 2");

            Ok(())
        });

        // Wait for the server to be ready
        let address = rx.await.expect("Failed to get server address");

        // Connect to the listener twice, sending some bytes and receiving some bytes from each
        let mut buf: [u8; 12] = [0; 12];

        let mut conn = TcpTransport::connect(&address)
            .await
            .expect("Conn 1 failed to connect");
        conn.write_all(b"hello server 1")
            .await
            .expect("Conn 1 failed to write");
        conn.read_exact(&mut buf)
            .await
            .expect("Conn 1 failed to read");
        assert_eq!(&buf, b"hello conn 1");

        let mut conn = TcpTransport::connect(&address)
            .await
            .expect("Conn 2 failed to connect");
        conn.write_all(b"hello server 2")
            .await
            .expect("Conn 2 failed to write");
        conn.read_exact(&mut buf)
            .await
            .expect("Conn 2 failed to read");
        assert_eq!(&buf, b"hello conn 2");

        // Verify that the task has completed by waiting on it
        let _ = task.await.expect("Listener task failed unexpectedly");
    }
}