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