Skip to content

Commit 60cc9d7

Browse files
Ayush1325Ayush1325
committed
std: sys: net: uefi: tcp4: Add timeout support
- Implement timeout support for read, write and connect. - A software implementation using Instant. Signed-off-by: Ayush Singh <[email protected]>
1 parent c83e217 commit 60cc9d7

File tree

3 files changed

+83
-25
lines changed

3 files changed

+83
-25
lines changed

‎library/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> {

‎library/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
}

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

Lines changed: 47 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,36 @@ impl Tcp4 {
167168
}
168169
}
169170

171+
/// # SAFETY
172+
/// Pointer to a token that has been issued by EFI_TCP4_PROTOCOL.Connect(),
173+
/// EFI_TCP4_PROTOCOL.Accept(), EFI_TCP4_PROTOCOL.Transmit() or EFI_TCP4_PROTOCOL.Receive().
174+
unsafe fn wait_or_cancel(
175+
&self,
176+
timeout: Option<Duration>,
177+
token: *mut tcp4::CompletionToken,
178+
) -> io::Result<()> {
179+
if !self.wait_for_flag(timeout) {
180+
let _ = unsafe { self.cancel(token) };
181+
return Err(io::Error::new(io::ErrorKind::TimedOut, "Operation Timed out"));
182+
}
183+
184+
Ok(())
185+
}
186+
187+
/// # SAFETY
188+
/// Pointer to a token that has been issued by EFI_TCP4_PROTOCOL.Connect(),
189+
/// EFI_TCP4_PROTOCOL.Accept(), EFI_TCP4_PROTOCOL.Transmit() or EFI_TCP4_PROTOCOL.Receive().
190+
unsafe fn cancel(&self, token: *mut tcp4::CompletionToken) -> io::Result<()> {
191+
let protocol = self.protocol.as_ptr();
192+
193+
let r = unsafe { ((*protocol).cancel)(protocol, token) };
194+
if r.is_error() {
195+
return Err(io::Error::from_raw_os_error(r.as_usize()));
196+
} else {
197+
Ok(())
198+
}
199+
}
200+
170201
unsafe fn create_evt(&self) -> io::Result<helpers::OwnedEvent> {
171202
self.flag.store(false, Ordering::Relaxed);
172203
helpers::OwnedEvent::new(
@@ -177,10 +208,19 @@ impl Tcp4 {
177208
)
178209
}
179210

180-
fn wait_for_flag(&self) {
211+
fn wait_for_flag(&self, timeout: Option<Duration>) -> bool {
212+
let start = Instant::now();
213+
181214
while !self.flag.load(Ordering::Relaxed) {
182215
let _ = self.poll();
216+
if let Some(t) = timeout {
217+
if Instant::now().duration_since(start) >= t {
218+
return false;
219+
}
220+
}
183221
}
222+
223+
true
184224
}
185225

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

0 commit comments

Comments
 (0)