librqbit_dualstack_sockets/
socket.rs

1use std::{
2    net::{IpAddr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
3    task::Poll,
4};
5
6use socket2::{Domain, Socket};
7use tracing::{debug, trace};
8
9use crate::{
10    Error,
11    addr::{ToV6Mapped, TryToV4},
12};
13
14#[derive(Clone, Copy, Debug)]
15pub enum SocketAddrKind {
16    V4(SocketAddrV4),
17    V6 {
18        addr: SocketAddrV6,
19        is_dualstack: bool,
20    },
21}
22
23impl SocketAddrKind {
24    fn as_socketaddr(&self) -> SocketAddr {
25        match *self {
26            SocketAddrKind::V4(addr) => SocketAddr::V4(addr),
27            SocketAddrKind::V6 { addr, .. } => SocketAddr::V6(addr),
28        }
29    }
30}
31
32pub struct MaybeDualstackSocket<S> {
33    socket: S,
34    addr_kind: SocketAddrKind,
35}
36
37impl<S> MaybeDualstackSocket<S> {
38    pub fn socket(&self) -> &S {
39        &self.socket
40    }
41
42    pub fn bind_addr(&self) -> SocketAddr {
43        self.addr_kind.as_socketaddr()
44    }
45
46    pub fn is_dualstack(&self) -> bool {
47        matches!(
48            self.addr_kind,
49            SocketAddrKind::V6 {
50                is_dualstack: true,
51                ..
52            }
53        )
54    }
55
56    fn convert_addr_for_send(&self, addr: SocketAddr) -> SocketAddr {
57        if self.is_dualstack() {
58            return SocketAddr::V6(addr.to_ipv6_mapped());
59        }
60        addr
61    }
62}
63
64impl MaybeDualstackSocket<Socket> {
65    fn bind(addr: SocketAddr, request_dualstack: bool, is_udp: bool) -> crate::Result<Self> {
66        let socket = Socket::new(
67            if addr.is_ipv6() {
68                Domain::IPV6
69            } else {
70                Domain::IPV4
71            },
72            if is_udp {
73                socket2::Type::DGRAM
74            } else {
75                socket2::Type::STREAM
76            },
77            Some(if is_udp {
78                socket2::Protocol::UDP
79            } else {
80                socket2::Protocol::TCP
81            }),
82        )
83        .map_err(Error::SocketNew)?;
84
85        let mut set_dualstack = false;
86
87        let addr_kind = match (request_dualstack, addr) {
88            (request_dualstack, SocketAddr::V6(addr))
89                if *addr.ip() == IpAddr::V6(Ipv6Addr::UNSPECIFIED) =>
90            {
91                let value = !request_dualstack;
92                trace!(?addr, only_v6 = value, "setting only_v6");
93                socket
94                    .set_only_v6(value)
95                    .map_err(|e| Error::OnlyV6 { value, source: e })?;
96                #[cfg(not(windows))] // socket.only_v6() panics on windows somehow
97                trace!(?addr, only_v6=?socket.only_v6());
98                set_dualstack = true;
99                SocketAddrKind::V6 {
100                    addr,
101                    is_dualstack: request_dualstack,
102                }
103            }
104            (_, SocketAddr::V6(addr)) => SocketAddrKind::V6 {
105                addr,
106                is_dualstack: false,
107            },
108            (_, SocketAddr::V4(addr)) => SocketAddrKind::V4(addr),
109        };
110
111        if !set_dualstack {
112            debug!(
113                ?addr,
114                "ignored dualstack request as it only applies to [::] address"
115            );
116        }
117
118        #[cfg(not(windows))]
119        {
120            if !is_udp {
121                socket
122                    .set_reuse_address(true)
123                    .map_err(Error::ReuseAddress)?;
124            }
125        }
126
127        socket.bind(&addr.into()).map_err(|e| {
128            trace!(?addr, "error binding: {e:#}");
129            Error::Bind(e)
130        })?;
131
132        let local_addr: SocketAddr = socket
133            .local_addr()
134            .map_err(Error::LocalAddr)?
135            .as_socket()
136            .ok_or(Error::AsSocket)?;
137
138        let addr_kind = match (addr_kind, local_addr) {
139            (SocketAddrKind::V4(..), SocketAddr::V4(received)) => SocketAddrKind::V4(received),
140            (SocketAddrKind::V6 { is_dualstack, .. }, SocketAddr::V6(received)) => {
141                SocketAddrKind::V6 {
142                    addr: received,
143                    is_dualstack,
144                }
145            }
146            _ => {
147                tracing::debug!(?local_addr, bind_addr=?addr, "mismatch between local_addr() and requested bind_addr");
148                return Err(Error::LocalBindAddrMismatch);
149            }
150        };
151
152        socket
153            .set_nonblocking(true)
154            .map_err(Error::SetNonblocking)?;
155
156        Ok(Self { socket, addr_kind })
157    }
158}
159
160impl MaybeDualstackSocket<tokio::net::TcpListener> {
161    pub fn bind_tcp(addr: SocketAddr, request_dualstack: bool) -> crate::Result<Self> {
162        let sock = MaybeDualstackSocket::bind(addr, request_dualstack, false)?;
163
164        debug!(addr=?sock.bind_addr(), requested_addr=?addr, dualstack = sock.is_dualstack(), "listening on TCP");
165        sock.socket().listen(1024).map_err(Error::Listen)?;
166
167        Ok(Self {
168            socket: tokio::net::TcpListener::from_std(std::net::TcpListener::from(sock.socket))
169                .map_err(Error::TokioFromStd)?,
170            addr_kind: sock.addr_kind,
171        })
172    }
173
174    pub async fn accept(&self) -> std::io::Result<(tokio::net::TcpStream, SocketAddr)> {
175        let (s, addr) = self.socket.accept().await?;
176        Ok((s, addr.try_to_ipv4()))
177    }
178}
179
180#[cfg(feature = "axum")]
181impl axum::serve::Listener for MaybeDualstackSocket<tokio::net::TcpListener> {
182    type Io = tokio::net::TcpStream;
183
184    type Addr = SocketAddr;
185
186    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
187        use backon::{ExponentialBuilder, Retryable};
188        (|| MaybeDualstackSocket::accept(self))
189            .retry(
190                ExponentialBuilder::new()
191                    .without_max_times()
192                    .with_max_delay(std::time::Duration::from_secs(5)),
193            )
194            .notify(|e, retry_in| tracing::trace!(?retry_in, "error accepting: {e:#}"))
195            .await
196            .unwrap()
197    }
198
199    fn local_addr(&self) -> tokio::io::Result<Self::Addr> {
200        Ok(self.bind_addr())
201    }
202}
203
204impl MaybeDualstackSocket<tokio::net::UdpSocket> {
205    pub fn bind_udp(addr: SocketAddr, request_dualstack: bool) -> crate::Result<Self> {
206        let sock = MaybeDualstackSocket::bind(addr, request_dualstack, true)?;
207
208        debug!(addr=?sock.bind_addr(), requested_addr=?addr, dualstack = sock.is_dualstack(), "listening on UDP");
209
210        Ok(Self {
211            socket: tokio::net::UdpSocket::from_std(std::net::UdpSocket::from(sock.socket))
212                .map_err(Error::TokioFromStd)?,
213            addr_kind: sock.addr_kind,
214        })
215    }
216
217    pub async fn recv_from(&self, buf: &mut [u8]) -> std::io::Result<(usize, SocketAddr)> {
218        let (size, addr) = self.socket.recv_from(buf).await?;
219        Ok((size, addr.try_to_ipv4()))
220    }
221
222    pub async fn send_to(&self, buf: &[u8], target: SocketAddr) -> std::io::Result<usize> {
223        let target = self.convert_addr_for_send(target);
224        self.socket.send_to(buf, target).await
225    }
226
227    pub fn poll_send_to(
228        &self,
229        cx: &mut std::task::Context<'_>,
230        buf: &[u8],
231        target: SocketAddr,
232    ) -> Poll<std::io::Result<usize>> {
233        let target = self.convert_addr_for_send(target);
234        self.socket.poll_send_to(cx, buf, target)
235    }
236}
librqbit_dualstack_sockets/socket.rs

librqbit_dualstack_sockets/
socket.rs
