From 311f0127de3adc209d670aaa3b952e10a18ce470 Mon Sep 17 00:00:00 2001
From: Jesse Braham <jessebraham@users.noreply.github.com>
Date: Wed, 31 Jul 2024 07:29:10 +0000
Subject: [PATCH] Implement `embedded_io::{ReadReady, WriteReady}` traits for
 `WifiStack` (#1882)

* implement `ReadReady` and `WriteReady` traits from `embedded-io` for Wi-Fi

* Allow non-static lifetime for `WifiStack`

* Update `CHANGELOG.md`

* clippy
---
 esp-wifi/CHANGELOG.md          |  6 ++--
 esp-wifi/src/wifi_interface.rs | 55 ++++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/esp-wifi/CHANGELOG.md b/esp-wifi/CHANGELOG.md
index 0ae57dbb337..e0ca0055c46 100644
--- a/esp-wifi/CHANGELOG.md
+++ b/esp-wifi/CHANGELOG.md
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Implement `embedded_io::{ReadReady, WriteReady}` traits for `WifiStack` (#1882)
+
 ### Changed
 
 ### Fixed
@@ -17,8 +19,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## 0.7.1 - 2024-07-17
 
-### Added
-
 ### Changed
 
 - Check no password is set when using `AuthMethod::None`(#1806)
@@ -27,8 +27,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Downgrade `embedded-svc` to 0.27.1 (#1820)
 
-### Removed
-
 ## 0.7.0 - 2024-07-15
 
 ### Added
diff --git a/esp-wifi/src/wifi_interface.rs b/esp-wifi/src/wifi_interface.rs
index 9c2910b77e5..7f8d33d0f7e 100644
--- a/esp-wifi/src/wifi_interface.rs
+++ b/esp-wifi/src/wifi_interface.rs
@@ -2,10 +2,6 @@
 
 use core::{borrow::BorrowMut, cell::RefCell, fmt::Display};
 
-#[cfg(feature = "tcp")]
-use embedded_io::ErrorType;
-#[cfg(feature = "tcp")]
-use embedded_io::{Read, Write};
 #[cfg(feature = "dhcpv4")]
 use smoltcp::socket::dhcpv4::Socket as Dhcpv4Socket;
 #[cfg(feature = "tcp")]
@@ -34,7 +30,7 @@ const LOCAL_PORT_MAX: u16 = 65535;
 ///
 /// Mostly a convenience wrapper for `smoltcp`
 pub struct WifiStack<'a, MODE: WifiDeviceMode> {
-    device: RefCell<WifiDevice<'static, MODE>>, // TODO allow non static lifetime
+    device: RefCell<WifiDevice<'a, MODE>>,
     network_interface: RefCell<Interface>,
     sockets: RefCell<SocketSet<'a>>,
     current_millis_fn: fn() -> u64,
@@ -53,7 +49,7 @@ pub struct WifiStack<'a, MODE: WifiDeviceMode> {
 impl<'a, MODE: WifiDeviceMode> WifiStack<'a, MODE> {
     pub fn new(
         network_interface: Interface,
-        device: WifiDevice<'static, MODE>, // TODO relax this lifetime requirement
+        device: WifiDevice<'a, MODE>,
         #[allow(unused_mut)] mut sockets: SocketSet<'a>,
         current_millis_fn: fn() -> u64,
     ) -> WifiStack<'a, MODE> {
@@ -709,12 +705,12 @@ impl embedded_io::Error for IoError {
 }
 
 #[cfg(feature = "tcp")]
-impl<'s, 'n: 's, MODE: WifiDeviceMode> ErrorType for Socket<'s, 'n, MODE> {
+impl<'s, 'n: 's, MODE: WifiDeviceMode> embedded_io::ErrorType for Socket<'s, 'n, MODE> {
     type Error = IoError;
 }
 
 #[cfg(feature = "tcp")]
-impl<'s, 'n: 's, MODE: WifiDeviceMode> Read for Socket<'s, 'n, MODE> {
+impl<'s, 'n: 's, MODE: WifiDeviceMode> embedded_io::Read for Socket<'s, 'n, MODE> {
     fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
         self.network.with_mut(|interface, device, sockets| {
             use smoltcp::socket::tcp::RecvError;
@@ -735,7 +731,25 @@ impl<'s, 'n: 's, MODE: WifiDeviceMode> Read for Socket<'s, 'n, MODE> {
 }
 
 #[cfg(feature = "tcp")]
-impl<'s, 'n: 's, MODE: WifiDeviceMode> Write for Socket<'s, 'n, MODE> {
+impl<'s, 'n: 's, MODE: WifiDeviceMode> embedded_io::ReadReady for Socket<'s, 'n, MODE> {
+    fn read_ready(&mut self) -> Result<bool, Self::Error> {
+        self.network.with_mut(|interface, device, sockets| {
+            use smoltcp::socket::tcp::RecvError;
+
+            interface.poll(timestamp(), device, sockets);
+            let socket = sockets.get_mut::<TcpSocket>(self.socket_handle);
+
+            match socket.peek(1) {
+                Ok(s) => Ok(!s.is_empty()),
+                Err(RecvError::Finished) => Err(IoError::SocketClosed),
+                Err(RecvError::InvalidState) => Err(IoError::TcpRecvError),
+            }
+        })
+    }
+}
+
+#[cfg(feature = "tcp")]
+impl<'s, 'n: 's, MODE: WifiDeviceMode> embedded_io::Write for Socket<'s, 'n, MODE> {
     fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
         loop {
             let (may_send, is_open, can_send) =
@@ -799,6 +813,29 @@ impl<'s, 'n: 's, MODE: WifiDeviceMode> Write for Socket<'s, 'n, MODE> {
     }
 }
 
+#[cfg(feature = "tcp")]
+impl<'s, 'n: 's, MODE: WifiDeviceMode> embedded_io::WriteReady for Socket<'s, 'n, MODE> {
+    fn write_ready(&mut self) -> Result<bool, Self::Error> {
+        let (may_send, is_open, can_send) = self.network.with_mut(|interface, device, sockets| {
+            interface.poll(
+                Instant::from_millis((self.network.current_millis_fn)() as i64),
+                device,
+                sockets,
+            );
+
+            let socket = sockets.get_mut::<TcpSocket>(self.socket_handle);
+
+            (socket.may_send(), socket.is_open(), socket.can_send())
+        });
+
+        if !is_open || !can_send {
+            return Err(IoError::SocketClosed);
+        }
+
+        Ok(may_send)
+    }
+}
+
 /// A UDP socket
 #[cfg(feature = "udp")]
 pub struct UdpSocket<'s, 'n: 's, MODE: WifiDeviceMode> {