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
128            .bind(&addr.into())
129            .map_err(|source| Error::Bind { addr, source })?;
130
131        let local_addr: SocketAddr = socket
132            .local_addr()
133            .map_err(Error::LocalAddr)?
134            .as_socket()
135            .ok_or(Error::AsSocket)?;
136
137        let addr_kind = match (addr_kind, local_addr) {
138            (SocketAddrKind::V4(..), SocketAddr::V4(received)) => SocketAddrKind::V4(received),
139            (SocketAddrKind::V6 { is_dualstack, .. }, SocketAddr::V6(received)) => {
140                SocketAddrKind::V6 {
141                    addr: received,
142                    is_dualstack,
143                }
144            }
145            _ => {
146                return Err(Error::LocalBindAddrMismatch {
147                    local_addr,
148                    bind_addr: addr,
149                });
150            }
151        };
152
153        socket
154            .set_nonblocking(true)
155            .map_err(Error::SetNonblocking)?;
156
157        Ok(Self { socket, addr_kind })
158    }
159}
160
161impl MaybeDualstackSocket<tokio::net::TcpListener> {
162    pub fn bind_tcp(addr: SocketAddr, request_dualstack: bool) -> crate::Result<Self> {
163        let sock = MaybeDualstackSocket::bind(addr, request_dualstack, false)?;
164
165        debug!(addr=?sock.bind_addr(), requested_addr=?addr, dualstack = sock.is_dualstack(), "listening on TCP");
166        sock.socket().listen(1024).map_err(Error::Listen)?;
167
168        Ok(Self {
169            socket: tokio::net::TcpListener::from_std(std::net::TcpListener::from(sock.socket))
170                .map_err(Error::TokioFromStd)?,
171            addr_kind: sock.addr_kind,
172        })
173    }
174
175    pub async fn accept(&self) -> std::io::Result<(tokio::net::TcpStream, SocketAddr)> {
176        let (s, addr) = self.socket.accept().await?;
177        Ok((s, addr.try_to_ipv4()))
178    }
179}
180
181impl MaybeDualstackSocket<tokio::net::UdpSocket> {
182    pub fn bind_udp(addr: SocketAddr, request_dualstack: bool) -> crate::Result<Self> {
183        let sock = MaybeDualstackSocket::bind(addr, request_dualstack, true)?;
184
185        debug!(addr=?sock.bind_addr(), requested_addr=?addr, dualstack = sock.is_dualstack(), "listening on UDP");
186
187        Ok(Self {
188            socket: tokio::net::UdpSocket::from_std(std::net::UdpSocket::from(sock.socket))
189                .map_err(Error::TokioFromStd)?,
190            addr_kind: sock.addr_kind,
191        })
192    }
193
194    pub async fn recv_from(&self, buf: &mut [u8]) -> std::io::Result<(usize, SocketAddr)> {
195        let (size, addr) = self.socket.recv_from(buf).await?;
196        Ok((size, addr.try_to_ipv4()))
197    }
198
199    pub async fn send_to(&self, buf: &[u8], target: SocketAddr) -> std::io::Result<usize> {
200        let target = self.convert_addr_for_send(target);
201        self.socket.send_to(buf, target).await
202    }
203
204    pub fn poll_send_to(
205        &self,
206        cx: &mut std::task::Context<'_>,
207        buf: &[u8],
208        target: SocketAddr,
209    ) -> Poll<std::io::Result<usize>> {
210        let target = self.convert_addr_for_send(target);
211        self.socket.poll_send_to(cx, buf, target)
212    }
213}
librqbit_dualstack_sockets/socket.rs

librqbit_dualstack_sockets/
socket.rs
