chore: unify ota and init states

3 files changed, 19 insertions, 18 deletions

diff --git a/embassy-net-esp-hosted/src/control.rs b/embassy-net-esp-hosted/src/control.rs
index 227db0d09..eb79593f6 100644
--- a/embassy-net-esp-hosted/src/control.rs
+++ b/embassy-net-esp-hosted/src/control.rs
@@ -177,7 +177,6 @@ impl<'a> Control<'a> {
     pub async fn ota_end(&mut self) -> Result<(), Error> {
         let req = proto::CtrlMsg_Req_OTAEnd {};
         ioctl!(self, ReqOtaEnd, RespOtaEnd, req, resp);
-        // Ensures that run loop awaits reset
         self.shared.ota_done();
         // Wait for re-init
         self.init().await?;
diff --git a/embassy-net-esp-hosted/src/ioctl.rs b/embassy-net-esp-hosted/src/ioctl.rs
index 8b9d582be..de0f867e8 100644
--- a/embassy-net-esp-hosted/src/ioctl.rs
+++ b/embassy-net-esp-hosted/src/ioctl.rs
@@ -23,18 +23,23 @@ pub struct Shared(RefCell<SharedInner>);
 
 struct SharedInner {
     ioctl: IoctlState,
-    is_init: bool,
-    is_ota: bool,
+    state: ControlState,
     control_waker: WakerRegistration,
     runner_waker: WakerRegistration,
 }
 
+#[derive(Clone, Copy)]
+pub(crate) enum ControlState {
+    Init,
+    Reboot,
+    Ready,
+}
+
 impl Shared {
     pub fn new() -> Self {
         Self(RefCell::new(SharedInner {
             ioctl: IoctlState::Done { resp_len: 0 },
-            is_init: false,
-            is_ota: false,
+            state: ControlState::Init,
             control_waker: WakerRegistration::new(),
             runner_waker: WakerRegistration::new(),
         }))
@@ -104,29 +109,27 @@ impl Shared {
     // ota
     pub fn ota_done(&self) {
         let mut this = self.0.borrow_mut();
-        this.is_ota = true;
-        this.is_init = false;
+        this.state = ControlState::Reboot;
     }
 
+    // // // // // // // // // // // // // // // // // // // //
+    //
     // check if ota is in progress
-    pub fn is_ota(&self) -> bool {
+    pub(crate) fn state(&self) -> ControlState {
         let this = self.0.borrow();
-        this.is_ota
+        this.state
     }
 
-    // // // // // // // // // // // // // // // // // // // //
-
     pub fn init_done(&self) {
         let mut this = self.0.borrow_mut();
-        this.is_init = true;
-        this.is_ota = false;
+        this.state = ControlState::Ready;
         this.control_waker.wake();
     }
 
     pub fn init_wait(&self) -> impl Future<Output = ()> + '_ {
         poll_fn(|cx| {
             let mut this = self.0.borrow_mut();
-            if this.is_init {
+            if let ControlState::Ready = this.state {
                 Poll::Ready(())
             } else {
                 this.control_waker.register(cx.waker());
diff --git a/embassy-net-esp-hosted/src/lib.rs b/embassy-net-esp-hosted/src/lib.rs
index f0cf04d04..2c7377281 100644
--- a/embassy-net-esp-hosted/src/lib.rs
+++ b/embassy-net-esp-hosted/src/lib.rs
@@ -234,13 +234,12 @@ where
                     tx_buf[..PayloadHeader::SIZE].fill(0);
                 }
                 Either4::Fourth(()) => {
-                    // Extend the deadline if OTA
-                    if self.shared.is_ota() {
+                    // Extend the deadline if initializing
+                    if let ioctl::ControlState::Reboot = self.shared.state() {
                         self.heartbeat_deadline = Instant::now() + HEARTBEAT_MAX_GAP;
                         continue;
-                    } else {
-                        panic!("heartbeat from esp32 stopped")
                     }
+                    panic!("heartbeat from esp32 stopped")
                 }
             }