kvarn_extensions/
connection.rs

1use kvarn::prelude::*;
2use std::net::{Ipv4Addr, SocketAddrV4};
3#[cfg(unix)]
4use tokio::net::UnixStream;
5use tokio::net::{TcpStream, UdpSocket};
6
7macro_rules! socket_addr_with_port {
8    ($($port:literal $(,)+)*) => {
9        &[
10            $(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, $port)),)*
11        ]
12    };
13}
14
15#[derive(Debug, Clone)]
16pub enum Connection {
17    Tcp(SocketAddr),
18    /// Keep in mind, this currently has a `60s` timeout.
19    /// Please use [`Self::UnixSocket`]s instead if you are on Unix.
20    Udp(SocketAddr),
21    #[cfg(unix)]
22    UnixSocket(String),
23}
24impl Connection {
25    pub async fn establish(self) -> io::Result<EstablishedConnection> {
26        match self {
27            Self::Tcp(addr) => TcpStream::connect(addr)
28                .await
29                .map(EstablishedConnection::Tcp),
30            Self::Udp(addr) => {
31                // Random ports to bind to.
32                let candidates = &socket_addr_with_port!(
33                    17448, 64567, 40022, 56654, 52027, 44328, 29973, 27919, 26513, 42327, 64855,
34                    5296, 52942, 43204, 15322, 13243,
35                )[..];
36                let socket = UdpSocket::bind(candidates).await?;
37                socket.connect(addr).await?;
38                Ok(EstablishedConnection::Udp(socket))
39            }
40            #[cfg(unix)]
41            Self::UnixSocket(path) => UnixStream::connect(path)
42                .await
43                .map(EstablishedConnection::UnixSocket),
44        }
45    }
46}
47
48#[derive(Debug)]
49pub enum EstablishedConnection {
50    Tcp(TcpStream),
51    Udp(UdpSocket),
52    #[cfg(unix)]
53    UnixSocket(UnixStream),
54}
55impl AsyncWrite for EstablishedConnection {
56    fn poll_write(
57        self: Pin<&mut Self>,
58        cx: &mut Context<'_>,
59        buf: &[u8],
60    ) -> Poll<Result<usize, io::Error>> {
61        match self.get_mut() {
62            Self::Tcp(s) => Pin::new(s).poll_write(cx, buf),
63            Self::Udp(s) => Pin::new(s).poll_send(cx, buf),
64            #[cfg(unix)]
65            Self::UnixSocket(s) => Pin::new(s).poll_write(cx, buf),
66        }
67    }
68    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
69        match self.get_mut() {
70            Self::Tcp(s) => Pin::new(s).poll_flush(cx),
71            Self::Udp(_) => Poll::Ready(Ok(())),
72            #[cfg(unix)]
73            Self::UnixSocket(s) => Pin::new(s).poll_flush(cx),
74        }
75    }
76    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
77        match self.get_mut() {
78            Self::Tcp(s) => Pin::new(s).poll_shutdown(cx),
79            Self::Udp(_) => Poll::Ready(Ok(())),
80            #[cfg(unix)]
81            Self::UnixSocket(s) => Pin::new(s).poll_shutdown(cx),
82        }
83    }
84}
85impl AsyncRead for EstablishedConnection {
86    fn poll_read(
87        self: Pin<&mut Self>,
88        cx: &mut Context<'_>,
89        buf: &mut ReadBuf<'_>,
90    ) -> Poll<io::Result<()>> {
91        match self.get_mut() {
92            Self::Tcp(s) => Pin::new(s).poll_read(cx, buf),
93            Self::Udp(s) => Pin::new(s).poll_recv(cx, buf),
94            #[cfg(unix)]
95            Self::UnixSocket(s) => Pin::new(s).poll_read(cx, buf),
96        }
97    }
98}
kvarn_extensions/connection.rs

kvarn_extensions/
connection.rs
