Skip to content

Commit d570e05

Browse files
matthiaskrgrmatthiaskrgr
authored
Rollup merge of rust-lang#143568 - Ayush1325:uefi-tcp4-timeout, r=tgross35
std: sys: net: uefi: tcp4: Add timeout support - Implement timeout support for read, write and connect. - A software implementation using Instant.
2 parents 1a42c32 + f9b215b commit d570e05

File tree

3 files changed

+97
-25
lines changed

3 files changed

+97
-25
lines changed

‎std/src/sys/net/connection/uefi/mod.rs‎

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,62 @@
11
use crate::fmt;
22
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
33
use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
4+
use crate::sync::{Arc, Mutex};
45
use crate::sys::unsupported;
56
use crate::time::Duration;
67

78
mod tcp;
89
pub(crate) mod tcp4;
910

10-
pub struct TcpStream(tcp::Tcp);
11+
pub struct TcpStream {
12+
inner: tcp::Tcp,
13+
read_timeout: Arc<Mutex<Option<Duration>>>,
14+
write_timeout: Arc<Mutex<Option<Duration>>>,
15+
}
1116

1217
impl TcpStream {
1318
pub fn connect(addr: io::Result<&SocketAddr>) -> io::Result<TcpStream> {
14-
tcp::Tcp::connect(addr?).map(Self)
19+
let inner = tcp::Tcp::connect(addr?, None)?;
20+
Ok(Self {
21+
inner,
22+
read_timeout: Arc::new(Mutex::new(None)),
23+
write_timeout: Arc::new(Mutex::new(None)),
24+
})
1525
}
1626

17-
pub fn connect_timeout(_: &SocketAddr, _: Duration) -> io::Result<TcpStream> {
18-
unsupported()
27+
pub fn connect_timeout(addr: &SocketAddr, timeout: Duration) -> io::Result<TcpStream> {
28+
let inner = tcp::Tcp::connect(addr, Some(timeout))?;
29+
Ok(Self {
30+
inner,
31+
read_timeout: Arc::new(Mutex::new(None)),
32+
write_timeout: Arc::new(Mutex::new(None)),
33+
})
1934
}
2035

21-
pub fn set_read_timeout(&self, _: Option<Duration>) -> io::Result<()> {
22-
unsupported()
36+
pub fn set_read_timeout(&self, t: Option<Duration>) -> io::Result<()> {
37+
self.read_timeout.set(t).unwrap();
38+
Ok(())
2339
}
2440

25-
pub fn set_write_timeout(&self, _: Option<Duration>) -> io::Result<()> {
26-
unsupported()
41+
pub fn set_write_timeout(&self, t: Option<Duration>) -> io::Result<()> {
42+
self.write_timeout.set(t).unwrap();
43+
Ok(())
2744
}
2845

2946
pub fn read_timeout(&self) -> io::Result<Option<Duration>> {
30-
unsupported()
47+
Ok(self.read_timeout.get_cloned().unwrap())
3148
}
3249

3350
pub fn write_timeout(&self) -> io::Result<Option<Duration>> {
34-
unsupported()
51+
Ok(self.write_timeout.get_cloned().unwrap())
3552
}
3653

3754
pub fn peek(&self, _: &mut [u8]) -> io::Result<usize> {
3855
unsupported()
3956
}
4057

4158
pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
42-
self.0.read(buf)
59+
self.inner.read(buf, self.read_timeout()?)
4360
}
4461

4562
pub fn read_buf(&self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
@@ -56,7 +73,7 @@ impl TcpStream {
5673
}
5774

5875
pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
59-
self.0.write(buf)
76+
self.inner.write(buf, self.write_timeout()?)
6077
}
6178

6279
pub fn write_vectored(&self, buf: &[IoSlice<'_>]) -> io::Result<usize> {

‎std/src/sys/net/connection/uefi/tcp.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,34 @@
11
use super::tcp4;
22
use crate::io;
33
use crate::net::SocketAddr;
4+
use crate::time::Duration;
45

56
pub(crate) enum Tcp {
67
V4(tcp4::Tcp4),
78
}
89

910
impl Tcp {
10-
pub(crate) fn connect(addr: &SocketAddr) -> io::Result<Self> {
11+
pub(crate) fn connect(addr: &SocketAddr, timeout: Option<Duration>) -> io::Result<Self> {
1112
match addr {
1213
SocketAddr::V4(x) => {
1314
let temp = tcp4::Tcp4::new()?;
1415
temp.configure(true, Some(x), None)?;
15-
temp.connect()?;
16+
temp.connect(timeout)?;
1617
Ok(Tcp::V4(temp))
1718
}
1819
SocketAddr::V6(_) => todo!(),
1920
}
2021
}
2122

22-
pub(crate) fn write(&self, buf: &[u8]) -> io::Result<usize> {
23+
pub(crate) fn write(&self, buf: &[u8], timeout: Option<Duration>) -> io::Result<usize> {
2324
match self {
24-
Self::V4(client) => client.write(buf),
25+
Self::V4(client) => client.write(buf, timeout),
2526
}
2627
}
2728

28-
pub(crate) fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
29+
pub(crate) fn read(&self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<usize> {
2930
match self {
30-
Self::V4(client) => client.read(buf),
31+
Self::V4(client) => client.read(buf, timeout),
3132
}
3233
}
3334
}

‎std/src/sys/net/connection/uefi/tcp4.rs‎

Lines changed: 61 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use crate::net::SocketAddrV4;
66
use crate::ptr::NonNull;
77
use crate::sync::atomic::{AtomicBool, Ordering};
88
use crate::sys::pal::helpers;
9+
use crate::time::{Duration, Instant};
910

1011
const TYPE_OF_SERVICE: u8 = 8;
1112
const TIME_TO_LIVE: u8 = 255;
@@ -66,7 +67,7 @@ impl Tcp4 {
6667
if r.is_error() { Err(crate::io::Error::from_raw_os_error(r.as_usize())) } else { Ok(()) }
6768
}
6869

69-
pub(crate) fn connect(&self) -> io::Result<()> {
70+
pub(crate) fn connect(&self, timeout: Option<Duration>) -> io::Result<()> {
7071
let evt = unsafe { self.create_evt() }?;
7172
let completion_token =
7273
tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS };
@@ -79,7 +80,7 @@ impl Tcp4 {
7980
return Err(io::Error::from_raw_os_error(r.as_usize()));
8081
}
8182

82-
self.wait_for_flag();
83+
unsafe { self.wait_or_cancel(timeout, &mut conn_token.completion_token) }?;
8384

8485
if completion_token.status.is_error() {
8586
Err(io::Error::from_raw_os_error(completion_token.status.as_usize()))
@@ -88,7 +89,7 @@ impl Tcp4 {
8889
}
8990
}
9091

91-
pub(crate) fn write(&self, buf: &[u8]) -> io::Result<usize> {
92+
pub(crate) fn write(&self, buf: &[u8], timeout: Option<Duration>) -> io::Result<usize> {
9293
let evt = unsafe { self.create_evt() }?;
9394
let completion_token =
9495
tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS };
@@ -119,7 +120,7 @@ impl Tcp4 {
119120
return Err(io::Error::from_raw_os_error(r.as_usize()));
120121
}
121122

122-
self.wait_for_flag();
123+
unsafe { self.wait_or_cancel(timeout, &mut token.completion_token) }?;
123124

124125
if completion_token.status.is_error() {
125126
Err(io::Error::from_raw_os_error(completion_token.status.as_usize()))
@@ -128,7 +129,7 @@ impl Tcp4 {
128129
}
129130
}
130131

131-
pub(crate) fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
132+
pub(crate) fn read(&self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<usize> {
132133
let evt = unsafe { self.create_evt() }?;
133134
let completion_token =
134135
tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS };
@@ -158,7 +159,7 @@ impl Tcp4 {
158159
return Err(io::Error::from_raw_os_error(r.as_usize()));
159160
}
160161

161-
self.wait_for_flag();
162+
unsafe { self.wait_or_cancel(timeout, &mut token.completion_token) }?;
162163

163164
if completion_token.status.is_error() {
164165
Err(io::Error::from_raw_os_error(completion_token.status.as_usize()))
@@ -167,6 +168,50 @@ impl Tcp4 {
167168
}
168169
}
169170

171+
/// Wait for an event to finish. This is checked by an atomic boolean that is supposed to be set
172+
/// to true in the event callback.
173+
///
174+
/// Optionally, allow specifying a timeout.
175+
///
176+
/// If a timeout is provided, the operation (specified by its `EFI_TCP4_COMPLETION_TOKEN`) is
177+
/// canceled and Error of kind TimedOut is returned.
178+
///
179+
/// # SAFETY
180+
///
181+
/// Pointer to a valid `EFI_TCP4_COMPLETION_TOKEN`
182+
unsafe fn wait_or_cancel(
183+
&self,
184+
timeout: Option<Duration>,
185+
token: *mut tcp4::CompletionToken,
186+
) -> io::Result<()> {
187+
if !self.wait_for_flag(timeout) {
188+
let _ = unsafe { self.cancel(token) };
189+
return Err(io::Error::new(io::ErrorKind::TimedOut, "Operation Timed out"));
190+
}
191+
192+
Ok(())
193+
}
194+
195+
/// Abort an asynchronous connection, listen, transmission or receive request.
196+
///
197+
/// If token is NULL, then all pending tokens issued by EFI_TCP4_PROTOCOL.Connect(),
198+
/// EFI_TCP4_PROTOCOL.Accept(), EFI_TCP4_PROTOCOL.Transmit() or EFI_TCP4_PROTOCOL.Receive() are
199+
/// aborted.
200+
///
201+
/// # SAFETY
202+
///
203+
/// Pointer to a valid `EFI_TCP4_COMPLETION_TOKEN` or NULL
204+
unsafe fn cancel(&self, token: *mut tcp4::CompletionToken) -> io::Result<()> {
205+
let protocol = self.protocol.as_ptr();
206+
207+
let r = unsafe { ((*protocol).cancel)(protocol, token) };
208+
if r.is_error() {
209+
return Err(io::Error::from_raw_os_error(r.as_usize()));
210+
} else {
211+
Ok(())
212+
}
213+
}
214+
170215
unsafe fn create_evt(&self) -> io::Result<helpers::OwnedEvent> {
171216
self.flag.store(false, Ordering::Relaxed);
172217
helpers::OwnedEvent::new(
@@ -177,10 +222,19 @@ impl Tcp4 {
177222
)
178223
}
179224

180-
fn wait_for_flag(&self) {
225+
fn wait_for_flag(&self, timeout: Option<Duration>) -> bool {
226+
let start = Instant::now();
227+
181228
while !self.flag.load(Ordering::Relaxed) {
182229
let _ = self.poll();
230+
if let Some(t) = timeout {
231+
if Instant::now().duration_since(start) >= t {
232+
return false;
233+
}
234+
}
183235
}
236+
237+
true
184238
}
185239

186240
fn poll(&self) -> io::Result<()> {

0 commit comments

Comments
 (0)