sqlx_core_guts/net/
socket.rs

1#![allow(dead_code)]
2
3use std::io;
4use std::path::Path;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use sqlx_rt::{AsyncRead, AsyncWrite, TcpStream};
9
10#[derive(Debug)]
11pub enum Socket {
12    Tcp(TcpStream),
13
14    #[cfg(unix)]
15    Unix(sqlx_rt::UnixStream),
16}
17
18impl Socket {
19    pub async fn connect_tcp(host: &str, port: u16) -> io::Result<Self> {
20        // Trim square brackets from host if it's an IPv6 address as the `url` crate doesn't do that.
21        TcpStream::connect((host.trim_matches(|c| c == '[' || c == ']'), port))
22            .await
23            .map(Socket::Tcp)
24    }
25
26    #[cfg(unix)]
27    pub async fn connect_uds(path: impl AsRef<Path>) -> io::Result<Self> {
28        sqlx_rt::UnixStream::connect(path.as_ref())
29            .await
30            .map(Socket::Unix)
31    }
32
33    #[cfg(not(unix))]
34    pub async fn connect_uds(_: impl AsRef<Path>) -> io::Result<Self> {
35        Err(io::Error::new(
36            io::ErrorKind::Other,
37            "Unix domain sockets are not supported outside Unix platforms.",
38        ))
39    }
40
41    pub async fn shutdown(&mut self) -> io::Result<()> {
42        #[cfg(feature = "_rt-async-std")]
43        {
44            use std::net::Shutdown;
45
46            match self {
47                Socket::Tcp(s) => s.shutdown(Shutdown::Both),
48
49                #[cfg(unix)]
50                Socket::Unix(s) => s.shutdown(Shutdown::Both),
51            }
52        }
53
54        #[cfg(any(feature = "_rt-actix", feature = "_rt-tokio"))]
55        {
56            use sqlx_rt::AsyncWriteExt;
57
58            match self {
59                Socket::Tcp(s) => s.shutdown().await,
60
61                #[cfg(unix)]
62                Socket::Unix(s) => s.shutdown().await,
63            }
64        }
65    }
66}
67
68impl AsyncRead for Socket {
69    fn poll_read(
70        mut self: Pin<&mut Self>,
71        cx: &mut Context<'_>,
72        buf: &mut super::PollReadBuf<'_>,
73    ) -> Poll<io::Result<super::PollReadOut>> {
74        match &mut *self {
75            Socket::Tcp(s) => Pin::new(s).poll_read(cx, buf),
76
77            #[cfg(unix)]
78            Socket::Unix(s) => Pin::new(s).poll_read(cx, buf),
79        }
80    }
81}
82
83impl AsyncWrite for Socket {
84    fn poll_write(
85        mut self: Pin<&mut Self>,
86        cx: &mut Context<'_>,
87        buf: &[u8],
88    ) -> Poll<io::Result<usize>> {
89        match &mut *self {
90            Socket::Tcp(s) => Pin::new(s).poll_write(cx, buf),
91
92            #[cfg(unix)]
93            Socket::Unix(s) => Pin::new(s).poll_write(cx, buf),
94        }
95    }
96
97    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
98        match &mut *self {
99            Socket::Tcp(s) => Pin::new(s).poll_flush(cx),
100
101            #[cfg(unix)]
102            Socket::Unix(s) => Pin::new(s).poll_flush(cx),
103        }
104    }
105
106    #[cfg(any(feature = "_rt-actix", feature = "_rt-tokio"))]
107    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
108        match &mut *self {
109            Socket::Tcp(s) => Pin::new(s).poll_shutdown(cx),
110
111            #[cfg(unix)]
112            Socket::Unix(s) => Pin::new(s).poll_shutdown(cx),
113        }
114    }
115
116    #[cfg(feature = "_rt-async-std")]
117    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
118        match &mut *self {
119            Socket::Tcp(s) => Pin::new(s).poll_close(cx),
120
121            #[cfg(unix)]
122            Socket::Unix(s) => Pin::new(s).poll_close(cx),
123        }
124    }
125}
sqlx_core_guts/net/socket.rs

sqlx_core_guts/net/
socket.rs
