|
| 1 | +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +use crate::msg; |
| 5 | +use crate::stream::socket::{fd::tcp, Flags, Socket}; |
| 6 | +use s2n_quic_core::inet::ExplicitCongestionNotification; |
| 7 | +use std::io::{self, Write}; |
| 8 | +use std::os::fd::AsRawFd; |
| 9 | +use std::pin::Pin; |
| 10 | +use std::task::Poll; |
| 11 | +use std::time::Duration; |
| 12 | +use std::{io::ErrorKind, net::TcpStream as StdTcpStream}; |
| 13 | +use tokio::{io::AsyncWrite as _, net::TcpStream as TokioTcpStream}; |
| 14 | + |
| 15 | +pub enum LazyBoundStream { |
| 16 | + Tokio(TokioTcpStream), |
| 17 | + Std(StdTcpStream), |
| 18 | + // needed for moving between the previous two while only having &mut access. |
| 19 | + TempEmpty, |
| 20 | +} |
| 21 | + |
| 22 | +impl LazyBoundStream { |
| 23 | + pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { |
| 24 | + match self { |
| 25 | + LazyBoundStream::Tokio(s) => s.set_nodelay(nodelay), |
| 26 | + LazyBoundStream::Std(s) => s.set_nodelay(nodelay), |
| 27 | + LazyBoundStream::TempEmpty => unreachable!(), |
| 28 | + } |
| 29 | + } |
| 30 | + |
| 31 | + pub fn set_linger(&self, linger: Option<Duration>) -> io::Result<()> { |
| 32 | + match self { |
| 33 | + LazyBoundStream::Tokio(s) => s.set_linger(linger), |
| 34 | + LazyBoundStream::Std(s) => { |
| 35 | + // Once it stabilizes we can switch to the std function |
| 36 | + // https://github.com/rust-lang/rust/issues/88494 |
| 37 | + let res = unsafe { |
| 38 | + libc::setsockopt( |
| 39 | + s.as_raw_fd(), |
| 40 | + libc::SOL_SOCKET, |
| 41 | + libc::SO_LINGER, |
| 42 | + &libc::linger { |
| 43 | + l_onoff: linger.is_some() as libc::c_int, |
| 44 | + l_linger: linger.unwrap_or_default().as_secs() as libc::c_int, |
| 45 | + } as *const _ as *const _, |
| 46 | + std::mem::size_of::<libc::linger>() as libc::socklen_t, |
| 47 | + ) |
| 48 | + }; |
| 49 | + if res != 0 { |
| 50 | + return Err(std::io::Error::last_os_error()); |
| 51 | + } |
| 52 | + |
| 53 | + Ok(()) |
| 54 | + } |
| 55 | + LazyBoundStream::TempEmpty => unreachable!(), |
| 56 | + } |
| 57 | + } |
| 58 | + |
| 59 | + pub fn into_std(self) -> io::Result<StdTcpStream> { |
| 60 | + match self { |
| 61 | + LazyBoundStream::Tokio(s) => s.into_std(), |
| 62 | + LazyBoundStream::Std(s) => Ok(s), |
| 63 | + LazyBoundStream::TempEmpty => unreachable!(), |
| 64 | + } |
| 65 | + } |
| 66 | + |
| 67 | + pub fn poll_write( |
| 68 | + &mut self, |
| 69 | + cx: &mut std::task::Context, |
| 70 | + buffer: &[u8], |
| 71 | + ) -> std::task::Poll<io::Result<usize>> { |
| 72 | + loop { |
| 73 | + match self { |
| 74 | + LazyBoundStream::Tokio(stream) => return Pin::new(stream).poll_write(cx, buffer), |
| 75 | + LazyBoundStream::Std(stream) => match stream.write(buffer) { |
| 76 | + Ok(v) => return Poll::Ready(Ok(v)), |
| 77 | + Err(e) => { |
| 78 | + if e.kind() == ErrorKind::WouldBlock { |
| 79 | + let LazyBoundStream::Std(stream) = |
| 80 | + std::mem::replace(self, LazyBoundStream::TempEmpty) |
| 81 | + else { |
| 82 | + unreachable!(); |
| 83 | + }; |
| 84 | + *self = LazyBoundStream::Tokio(TokioTcpStream::from_std(stream)?); |
| 85 | + } else { |
| 86 | + return Poll::Ready(Err(e)); |
| 87 | + } |
| 88 | + } |
| 89 | + }, |
| 90 | + LazyBoundStream::TempEmpty => unreachable!(), |
| 91 | + } |
| 92 | + } |
| 93 | + } |
| 94 | + |
| 95 | + pub fn poll_recv_buffer( |
| 96 | + &mut self, |
| 97 | + cx: &mut std::task::Context, |
| 98 | + buffer: &mut msg::recv::Message, |
| 99 | + ) -> std::task::Poll<io::Result<usize>> { |
| 100 | + loop { |
| 101 | + match self { |
| 102 | + LazyBoundStream::Tokio(stream) => { |
| 103 | + return Pin::new(stream).poll_recv_buffer(cx, buffer) |
| 104 | + } |
| 105 | + LazyBoundStream::Std(stream) => { |
| 106 | + let res = buffer.recv_with(|_addr, cmsg, buffer| { |
| 107 | + loop { |
| 108 | + let flags = Flags::default(); |
| 109 | + let res = tcp::recv(&*stream, buffer, flags); |
| 110 | + |
| 111 | + match res { |
| 112 | + Ok(len) => { |
| 113 | + // we don't need ECN markings from TCP since it handles that logic for us |
| 114 | + cmsg.set_ecn(ExplicitCongestionNotification::NotEct); |
| 115 | + |
| 116 | + // TCP doesn't have segments so just set it to 0 (which will indicate a single |
| 117 | + // stream of bytes) |
| 118 | + cmsg.set_segment_len(0); |
| 119 | + |
| 120 | + return Ok(len); |
| 121 | + } |
| 122 | + Err(ref e) if e.kind() == io::ErrorKind::Interrupted => { |
| 123 | + // try the operation again if we were interrupted |
| 124 | + continue; |
| 125 | + } |
| 126 | + Err(err) => return Err(err), |
| 127 | + } |
| 128 | + } |
| 129 | + }); |
| 130 | + match res { |
| 131 | + Ok(v) => return Poll::Ready(Ok(v)), |
| 132 | + Err(e) => { |
| 133 | + if e.kind() == ErrorKind::WouldBlock { |
| 134 | + let LazyBoundStream::Std(stream) = |
| 135 | + std::mem::replace(self, LazyBoundStream::TempEmpty) |
| 136 | + else { |
| 137 | + unreachable!(); |
| 138 | + }; |
| 139 | + *self = LazyBoundStream::Tokio(TokioTcpStream::from_std(stream)?); |
| 140 | + } else { |
| 141 | + return Poll::Ready(Err(e)); |
| 142 | + } |
| 143 | + } |
| 144 | + } |
| 145 | + } |
| 146 | + LazyBoundStream::TempEmpty => unreachable!(), |
| 147 | + } |
| 148 | + } |
| 149 | + } |
| 150 | +} |
0 commit comments