lib.rs (21684B)
1 //! [![git]](https://git.philomathiclife.com/tokio_dual_stack/log.html) [![crates-io]](https://crates.io/crates/tokio_dual_stack) [![docs-rs]](crate) 2 //! 3 //! [git]: https://git.philomathiclife.com/git_badge.svg 4 //! [crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust 5 //! [docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs 6 //! 7 //! `tokio_dual_stack` is a library that adds a "dual-stack" [`TcpListener`]. 8 //! 9 //! ## Why is this useful? 10 //! 11 //! Only certain platforms offer the ability for one socket to handle both IPv6 and IPv4 requests 12 //! (e.g., OpenBSD does not). For the platforms that do, it is often dependent on runtime configuration 13 //! (e.g., [`IPV6_V6ONLY`](https://www.man7.org/linux/man-pages/man7/ipv6.7.html)). Additionally those platforms 14 //! that support it often require the "wildcard" IPv6 address to be used (i.e., `::`) which has the unfortunate 15 //! consequence of preventing other services from using the same protocol port. 16 //! 17 //! There are a few ways to work around this issue. One is to deploy the same service twice: one that uses 18 //! an IPv6 socket and the other that uses an IPv4 socket. This can complicate deployments (e.g., the application 19 //! may not have been written with the expectation that multiple deployments could be running at the same time) in 20 //! addition to using more resources. Another is for the application to manually handle each socket (e.g., 21 //! [`select`](https://docs.rs/tokio/latest/tokio/macro.select.html)/[`join`](https://docs.rs/tokio/latest/tokio/macro.join.html) 22 //! each [`TcpListener::accept`]). 23 //! 24 //! [`DualStackTcpListener`] chooses an implementation similar to what the equivalent `select` would do while 25 //! also ensuring that one socket does not "starve" another by ensuring each socket is fairly given an opportunity 26 //! to `TcpListener::accept` a connection. This has the nice benefit of having a similar API to what a single 27 //! `TcpListener` would have as well as having similar performance to a socket that does handle both IPv6 and 28 //! IPv4 requests. 29 #![expect( 30 clippy::doc_paragraphs_missing_punctuation, 31 reason = "false positive for crate documentation having image links" 32 )] 33 #![expect( 34 clippy::std_instead_of_core, 35 reason = "false positive until core::io::ErrorKind is stable" 36 )] 37 #![cfg_attr( 38 test, 39 expect(dead_code_pub_in_binary, reason = "ignore for unit tests") 40 )] 41 use core::{ 42 net::{SocketAddr, SocketAddrV4, SocketAddrV6}, 43 pin::Pin, 44 sync::atomic::{AtomicBool, Ordering}, 45 task::{Context, Poll}, 46 }; 47 use pin_project_lite::pin_project; 48 use std::io::{Error, ErrorKind, Result}; 49 pub use tokio; 50 use tokio::net::{self, TcpListener, TcpSocket, TcpStream, ToSocketAddrs}; 51 /// Prevents [`Sealed`] from being publicly implementable. 52 mod private; 53 use private::Sealed; 54 /// TCP "listener". 55 /// 56 /// This `trait` is sealed and cannot be implemented for types outside of `tokio_dual_stack`. 57 /// 58 /// This exists primarily as a way to define type constructors or polymorphic functions 59 /// that can user either a [`TcpListener`] or [`DualStackTcpListener`]. 60 /// 61 /// # Examples 62 /// 63 /// ```no_run 64 /// # use tokio_dual_stack::Tcp; 65 /// async fn main_loop<T: Tcp>(listener: T) -> ! { 66 /// loop { 67 /// match listener.accept().await { 68 /// Ok((_, socket)) => println!("Client socket: {socket}"), 69 /// Err(e) => println!("TCP connection failure: {e}"), 70 /// } 71 /// } 72 /// } 73 /// ``` 74 pub trait Tcp: Sealed + Sized { 75 /// Creates a new TCP listener, which will be bound to the specified address(es). 76 /// 77 /// The returned listener is ready for accepting connections. 78 /// 79 /// Binding with a port number of 0 will request that the OS assigns a port to this listener. 80 /// The port allocated can be queried via the `local_addr` method. 81 /// 82 /// The address type can be any implementor of the [`ToSocketAddrs`] trait. If `addr` yields 83 /// multiple addresses, bind will be attempted with each of the addresses until one succeeds 84 /// and returns the listener. If none of the addresses succeed in creating a listener, the 85 /// error returned from the last attempt (the last address) is returned. 86 /// 87 /// This function sets the `SO_REUSEADDR` option on the socket. 88 /// 89 /// # Examples 90 /// 91 /// ```no_run 92 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 93 /// # use std::io::Result; 94 /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _}; 95 /// #[tokio::main(flavor = "current_thread")] 96 /// async fn main() -> Result<()> { 97 /// let listener = DualStackTcpListener::bind( 98 /// [ 99 /// SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)), 100 /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)), 101 /// ] 102 /// .as_slice(), 103 /// ) 104 /// .await?; 105 /// Ok(()) 106 /// } 107 /// ``` 108 fn bind<A: ToSocketAddrs>(addr: A) -> impl Future<Output = Result<Self>>; 109 /// Accepts a new incoming connection from this listener. 110 /// 111 /// This function will yield once a new TCP connection is established. When established, 112 /// the corresponding `TcpStream` and the remote peer’s address will be returned. 113 /// 114 /// # Cancel safety 115 /// 116 /// This method is cancel safe. If the method is used as the event in a 117 /// [`tokio::select!`](https://docs.rs/tokio/latest/tokio/macro.select.html) 118 /// statement and some other branch completes first, then it is guaranteed that no new 119 /// connections were accepted by this method. 120 /// 121 /// # Examples 122 /// 123 /// ```no_run 124 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 125 /// # use std::io::Result; 126 /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _}; 127 /// #[tokio::main(flavor = "current_thread")] 128 /// async fn main() -> Result<()> { 129 /// match DualStackTcpListener::bind( 130 /// [ 131 /// SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)), 132 /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)), 133 /// ] 134 /// .as_slice(), 135 /// ) 136 /// .await?.accept().await { 137 /// Ok((_, addr)) => println!("new client: {addr}"), 138 /// Err(e) => println!("couldn't get client: {e}"), 139 /// } 140 /// Ok(()) 141 /// } 142 /// ``` 143 fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync; 144 /// Polls to accept a new incoming connection to this listener. 145 /// 146 /// If there is no connection to accept, `Poll::Pending` is returned and the current task will be notified by 147 /// a waker. Note that on multiple calls to `poll_accept`, only the `Waker` from the `Context` passed to the 148 /// most recent call is scheduled to receive a wakeup. 149 fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>>; 150 } 151 impl Sealed for TcpListener {} 152 impl Tcp for TcpListener { 153 #[inline] 154 fn bind<A: ToSocketAddrs>(addr: A) -> impl Future<Output = Result<Self>> { 155 Self::bind(addr) 156 } 157 #[inline] 158 fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync { 159 self.accept() 160 } 161 #[inline] 162 fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>> { 163 self.poll_accept(cx) 164 } 165 } 166 /// "Dual-stack" TCP listener. 167 /// 168 /// IPv6 and IPv4 TCP listener. 169 #[derive(Debug)] 170 pub struct DualStackTcpListener { 171 /// IPv6 TCP listener. 172 ip6: TcpListener, 173 /// IPv4 TCP listener. 174 ip4: TcpListener, 175 /// `true` iff [`Self::ip6::accept`] should be `poll`ed first; otherwise [`Self::ip4::accept`] is `poll`ed 176 /// first. 177 /// 178 /// This exists to prevent one IP version from "starving" another. Each time [`Self::accept`] or 179 /// [`Self::poll_accept`] is called, it's overwritten with the opposite `bool`. 180 /// 181 /// Note we could make this a `core::cell::Cell`; but for maximal flexibility and consistency with `TcpListener`, 182 /// we use an `AtomicBool`. This among other things means `DualStackTcpListener` will implement `Sync`. 183 ip6_first: AtomicBool, 184 } 185 impl DualStackTcpListener { 186 /// Creates `Self` using the [`TcpListener`]s returned from [`TcpSocket::listen`]. 187 /// 188 /// [`Self::bind`] is useful when the behavior of [`TcpListener::bind`] is sufficient; however if the underlying 189 /// `TcpSocket`s need to be configured differently, then one must call this function instead. 190 /// 191 /// # Errors 192 /// 193 /// Errors iff [`TcpSocket::local_addr`] does for either socket, the underlying sockets use the same IP version, 194 /// or [`TcpSocket::listen`] errors for either socket. 195 /// 196 /// Note on Windows-based platforms `TcpSocket::local_addr` will error if [`TcpSocket::bind`] was not called. 197 /// 198 /// # Examples 199 /// 200 /// ```no_run 201 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 202 /// # use std::io::Result; 203 /// # use tokio_dual_stack::DualStackTcpListener; 204 /// # use tokio::net::TcpSocket; 205 /// #[tokio::main(flavor = "current_thread")] 206 /// async fn main() -> Result<()> { 207 /// let ip6 = TcpSocket::new_v6()?; 208 /// ip6.bind(SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)))?; 209 /// let ip4 = TcpSocket::new_v4()?; 210 /// ip4.bind(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)))?; 211 /// let listener = DualStackTcpListener::from_sockets((ip6, 1024), (ip4, 1024))?; 212 /// Ok(()) 213 /// } 214 /// ``` 215 #[inline] 216 pub fn from_sockets( 217 (socket_1, backlog_1): (TcpSocket, u32), 218 (socket_2, backlog_2): (TcpSocket, u32), 219 ) -> Result<Self> { 220 socket_1.local_addr().and_then(|sock| { 221 socket_2.local_addr().and_then(|sock_2| { 222 if sock.is_ipv6() { 223 if sock_2.is_ipv4() { 224 socket_1.listen(backlog_1).and_then(|ip6| { 225 socket_2.listen(backlog_2).map(|ip4| Self { 226 ip6, 227 ip4, 228 ip6_first: AtomicBool::new(true), 229 }) 230 }) 231 } else { 232 Err(Error::new( 233 ErrorKind::InvalidData, 234 "TcpSockets are the same IP version", 235 )) 236 } 237 } else if sock_2.is_ipv6() { 238 socket_1.listen(backlog_1).and_then(|ip4| { 239 socket_2.listen(backlog_2).map(|ip6| Self { 240 ip6, 241 ip4, 242 ip6_first: AtomicBool::new(true), 243 }) 244 }) 245 } else { 246 Err(Error::new( 247 ErrorKind::InvalidData, 248 "TcpSockets are the same IP version", 249 )) 250 } 251 }) 252 }) 253 } 254 /// Returns the local address of each socket that the listeners are bound to. 255 /// 256 /// This can be useful, for example, when binding to port 0 to figure out which port was actually bound. 257 /// 258 /// # Errors 259 /// 260 /// Errors iff [`TcpListener::local_addr`] does for either listener. 261 /// 262 /// # Examples 263 /// 264 /// ```no_run 265 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 266 /// # use std::io::Result; 267 /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _}; 268 /// #[tokio::main(flavor = "current_thread")] 269 /// async fn main() -> Result<()> { 270 /// let ip6 = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0); 271 /// let ip4 = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080); 272 /// assert_eq!( 273 /// DualStackTcpListener::bind([SocketAddr::V6(ip6), SocketAddr::V4(ip4)].as_slice()) 274 /// .await? 275 /// .local_addr()?, 276 /// (ip6, ip4) 277 /// ); 278 /// Ok(()) 279 /// } 280 /// ``` 281 #[expect(clippy::unreachable, reason = "we want to crash when there is a bug")] 282 #[inline] 283 pub fn local_addr(&self) -> Result<(SocketAddrV6, SocketAddrV4)> { 284 self.ip6.local_addr().and_then(|ip6| { 285 self.ip4.local_addr().map(|ip4| { 286 ( 287 if let SocketAddr::V6(sock6) = ip6 { 288 sock6 289 } else { 290 unreachable!("there is a bug in DualStackTcpListener::bind") 291 }, 292 if let SocketAddr::V4(sock4) = ip4 { 293 sock4 294 } else { 295 unreachable!("there is a bug in DualStackTcpListener::bind") 296 }, 297 ) 298 }) 299 }) 300 } 301 /// Sets the value for the `IP_TTL` option on both sockets. 302 /// 303 /// This value sets the time-to-live field that is used in every packet sent from each socket. 304 /// `ttl_ip6` is the `IP_TTL` value for the IPv6 socket and `ttl_ip4` is the `IP_TTL` value for the 305 /// IPv4 socket. 306 /// 307 /// # Errors 308 /// 309 /// Errors iff [`TcpListener::set_ttl`] does for either listener. 310 /// 311 /// # Examples 312 /// 313 /// ```no_run 314 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 315 /// # use std::io::Result; 316 /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _}; 317 /// #[tokio::main(flavor = "current_thread")] 318 /// async fn main() -> Result<()> { 319 /// DualStackTcpListener::bind( 320 /// [ 321 /// SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)), 322 /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)), 323 /// ] 324 /// .as_slice(), 325 /// ) 326 /// .await?.set_ttl(100, 100).expect("could not set TTL"); 327 /// Ok(()) 328 /// } 329 /// ``` 330 #[inline] 331 pub fn set_ttl(&self, ttl_ip6: u32, ttl_ip4: u32) -> Result<()> { 332 self.ip6 333 .set_ttl(ttl_ip6) 334 .and_then(|()| self.ip4.set_ttl(ttl_ip4)) 335 } 336 /// Gets the values of the `IP_TTL` option for both sockets. 337 /// 338 /// The first `u32` represents the `IP_TTL` value for the IPv6 socket and the second `u32` is the 339 /// `IP_TTL` value for the IPv4 socket. For more information about this option, see [`Self::set_ttl`]. 340 /// 341 /// # Errors 342 /// 343 /// Errors iff [`TcpListener::ttl`] does for either listener. 344 /// 345 /// # Examples 346 /// 347 /// ```no_run 348 /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; 349 /// # use std::io::Result; 350 /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _}; 351 /// #[tokio::main(flavor = "current_thread")] 352 /// async fn main() -> Result<()> { 353 /// let listener = DualStackTcpListener::bind( 354 /// [ 355 /// SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)), 356 /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)), 357 /// ] 358 /// .as_slice(), 359 /// ) 360 /// .await?; 361 /// listener.set_ttl(100, 100).expect("could not set TTL"); 362 /// assert_eq!(listener.ttl()?, (100, 100)); 363 /// Ok(()) 364 /// } 365 /// ``` 366 #[inline] 367 pub fn ttl(&self) -> Result<(u32, u32)> { 368 self.ip6 369 .ttl() 370 .and_then(|ip6| self.ip4.ttl().map(|ip4| (ip6, ip4))) 371 } 372 } 373 pin_project! { 374 /// `Future` returned by [`DualStackTcpListener::accept]`. 375 struct AcceptFut< 376 F: Future<Output = Result<(TcpStream, SocketAddr)>>, 377 F2: Future<Output = Result<(TcpStream, SocketAddr)>>, 378 > { 379 // Accept future for one `TcpListener`. 380 #[pin] 381 fut_1: F, 382 // Accept future for the other `TcpListener`. 383 #[pin] 384 fut_2: F2, 385 } 386 } 387 impl< 388 F: Future<Output = Result<(TcpStream, SocketAddr)>>, 389 F2: Future<Output = Result<(TcpStream, SocketAddr)>>, 390 > Future for AcceptFut<F, F2> 391 { 392 type Output = Result<(TcpStream, SocketAddr)>; 393 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 394 let this = self.project(); 395 // Note we defer errors caused from polling a completed `Future` to the contained `tokio` `Future`s. 396 // The only time this `Future` can be polled after completion without an error (due to `tokio`) is 397 // if `fut_2` completes first, `self` is polled, then `fut_1` completes. We don't actually care 398 // that this happens since the correctness of the code is still fine. 399 // This means any bugs that could occur from polling this `Future` after completion are dependency-based 400 // bugs where the correct solution is to fix the bugs in `tokio`. 401 match this.fut_1.poll(cx) { 402 Poll::Ready(res) => Poll::Ready(res), 403 Poll::Pending => this.fut_2.poll(cx), 404 } 405 } 406 } 407 impl Sealed for DualStackTcpListener {} 408 impl Tcp for DualStackTcpListener { 409 #[inline] 410 async fn bind<A: ToSocketAddrs>(addr: A) -> Result<Self> { 411 match net::lookup_host(addr).await { 412 Ok(socks) => { 413 let mut last_err = None; 414 let mut ip6_opt = None; 415 let mut ip4_opt = None; 416 for sock in socks { 417 match ip6_opt { 418 None => match ip4_opt { 419 None => { 420 let is_ip6 = sock.is_ipv6(); 421 match TcpListener::bind(sock).await { 422 Ok(ip) => { 423 if is_ip6 { 424 ip6_opt = Some(ip); 425 } else { 426 ip4_opt = Some(ip); 427 } 428 } 429 Err(err) => last_err = Some(err), 430 } 431 } 432 Some(ip4) => { 433 if sock.is_ipv6() { 434 match TcpListener::bind(sock).await { 435 Ok(ip6) => { 436 return Ok(Self { 437 ip6, 438 ip4, 439 ip6_first: AtomicBool::new(true), 440 }); 441 } 442 Err(err) => last_err = Some(err), 443 } 444 } 445 ip4_opt = Some(ip4); 446 } 447 }, 448 Some(ip6) => { 449 if sock.is_ipv4() { 450 match TcpListener::bind(sock).await { 451 Ok(ip4) => { 452 return Ok(Self { 453 ip6, 454 ip4, 455 ip6_first: AtomicBool::new(true), 456 }); 457 } 458 Err(err) => last_err = Some(err), 459 } 460 } 461 ip6_opt = Some(ip6); 462 } 463 } 464 } 465 Err(last_err.unwrap_or_else(|| { 466 Error::new( 467 ErrorKind::InvalidInput, 468 "could not resolve to an IPv6 and IPv4 address", 469 ) 470 })) 471 } 472 Err(err) => Err(err), 473 } 474 } 475 #[inline] 476 fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync { 477 // The correctness of code does not depend on `self.ip6_first`; therefore 478 // we elect for the most performant `Ordering`. 479 if self.ip6_first.swap(false, Ordering::Relaxed) { 480 AcceptFut { 481 fut_1: self.ip6.accept(), 482 fut_2: self.ip4.accept(), 483 } 484 } else { 485 // The correctness of code does not depend on `self.ip6_first`; therefore 486 // we elect for the most performant `Ordering`. 487 self.ip6_first.store(true, Ordering::Relaxed); 488 AcceptFut { 489 fut_1: self.ip4.accept(), 490 fut_2: self.ip6.accept(), 491 } 492 } 493 } 494 #[inline] 495 fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>> { 496 // The correctness of code does not depend on `self.ip6_first`; therefore 497 // we elect for the most performant `Ordering`. 498 if self.ip6_first.swap(false, Ordering::Relaxed) { 499 self.ip6.poll_accept(cx) 500 } else { 501 // The correctness of code does not depend on `self.ip6_first`; therefore 502 // we elect for the most performant `Ordering`. 503 self.ip6_first.store(true, Ordering::Relaxed); 504 self.ip4.poll_accept(cx) 505 } 506 } 507 }