4 files changed, 73 insertions, 7 deletions

diff --git a/embassy-net-esp-hosted/src/control.rs b/embassy-net-esp-hosted/src/control.rs
index 38ec648b4..d96a62daf 100644
--- a/embassy-net-esp-hosted/src/control.rs
+++ b/embassy-net-esp-hosted/src/control.rs
@@ -24,6 +24,11 @@ pub struct Control<'a> {
     shared: &'a Shared,
 }
 
+/// Handle for managing firmware update.
+pub struct UpdateControl<'a, 'd> {
+    control: &'a mut Control<'d>,
+}
+
 /// WiFi mode.
 #[allow(unused)]
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
@@ -146,6 +151,13 @@ impl<'a> Control<'a> {
         Ok(())
     }
 
+    /// Initiate a firmware update.
+    pub async fn update(&mut self) -> Result<UpdateControl<'_, 'a>, Error> {
+        let req = proto::CtrlMsg_Req_OTABegin {};
+        ioctl!(self, ReqOtaBegin, RespOtaBegin, req, resp);
+        Ok(UpdateControl { control: self })
+    }
+
     /// duration in seconds, clamped to [10, 3600]
     async fn set_heartbeat(&mut self, duration: u32) -> Result<(), Error> {
         let req = proto::CtrlMsg_Req_ConfigHeartbeat {
@@ -175,7 +187,8 @@ impl<'a> Control<'a> {
     async fn ioctl(&mut self, msg: &mut CtrlMsg) -> Result<(), Error> {
         debug!("ioctl req: {:?}", &msg);
 
-        let mut buf = [0u8; 128];
+        // Theoretical max overhead is 29 bytes. Biggest message is OTA write with 256 bytes.
+        let mut buf = [0u8; 256 + 29];
         let buf_len = buf.len();
 
         let mut encoder = PbEncoder::new(&mut buf[..]);
@@ -216,6 +229,37 @@ impl<'a> Control<'a> {
     }
 }
 
+impl<'a, 'd> UpdateControl<'a, 'd> {
+    /// Write slice of firmware to a device.
+    ///
+    /// The slice is split into chunks that can be sent across
+    /// the ioctl protocol to the wifi adapter.
+    pub async fn write(&mut self, data: &[u8]) -> Result<(), Error> {
+        let this = &mut self.control;
+        for chunk in data.chunks(256) {
+            let req = proto::CtrlMsg_Req_OTAWrite {
+                ota_data: heapless::Vec::from_slice(chunk).unwrap(),
+            };
+            ioctl!(this, ReqOtaWrite, RespOtaWrite, req, resp);
+        }
+        Ok(())
+    }
+
+    /// End the OTA session.
+    ///
+    /// NOTE: Will reset the wifi adapter after 5 seconds.
+    pub async fn finish(self) -> Result<(), Error> {
+        let this = self.control;
+        let req = proto::CtrlMsg_Req_OTAEnd {};
+        ioctl!(this, ReqOtaEnd, RespOtaEnd, req, resp);
+        // Ensures that run loop awaits reset
+        this.shared.ota_done();
+        // Wait for re-init
+        this.init().await?;
+        Ok(())
+    }
+}
+
 // WHY IS THIS A STRING? WHYYYY
 fn parse_mac(mac: &str) -> Result<[u8; 6], Error> {
     fn nibble_from_hex(b: u8) -> Result<u8, Error> {
diff --git a/embassy-net-esp-hosted/src/ioctl.rs b/embassy-net-esp-hosted/src/ioctl.rs
index a516f80c7..7f462d528 100644
--- a/embassy-net-esp-hosted/src/ioctl.rs
+++ b/embassy-net-esp-hosted/src/ioctl.rs
@@ -24,6 +24,7 @@ pub struct Shared(RefCell<SharedInner>);
 struct SharedInner {
     ioctl: IoctlState,
     is_init: bool,
+    is_ota: bool,
     control_waker: WakerRegistration,
     runner_waker: WakerRegistration,
 }
@@ -33,6 +34,7 @@ impl Shared {
         Self(RefCell::new(SharedInner {
             ioctl: IoctlState::Done { resp_len: 0 },
             is_init: false,
+            is_ota: false,
             control_waker: WakerRegistration::new(),
             runner_waker: WakerRegistration::new(),
         }))
@@ -99,11 +101,26 @@ impl Shared {
         }
     }
 
+    // ota
+    pub fn ota_done(&self) {
+        let mut this = self.0.borrow_mut();
+        this.is_ota = true;
+        this.is_init = false;
+        this.runner_waker.wake();
+    }
+
+    // check if ota is in progress
+    pub fn is_ota(&self) -> bool {
+        let this = self.0.borrow();
+        this.is_ota
+    }
+
     // // // // // // // // // // // // // // // // // // // //
 
     pub fn init_done(&self) {
         let mut this = self.0.borrow_mut();
         this.is_init = true;
+        this.is_ota = false;
         this.control_waker.wake();
     }
 
diff --git a/embassy-net-esp-hosted/src/lib.rs b/embassy-net-esp-hosted/src/lib.rs
index d882af8cf..7236e73e8 100644
--- a/embassy-net-esp-hosted/src/lib.rs
+++ b/embassy-net-esp-hosted/src/lib.rs
@@ -234,6 +234,11 @@ where
                     tx_buf[..PayloadHeader::SIZE].fill(0);
                 }
                 Either4::Fourth(()) => {
+                    // Extend the deadline if OTA
+                    if self.shared.is_ota() {
+                        self.heartbeat_deadline = Instant::now() + HEARTBEAT_MAX_GAP;
+                        continue;
+                    }
                     panic!("heartbeat from esp32 stopped")
                 }
             }
diff --git a/embassy-net-esp-hosted/src/proto.rs b/embassy-net-esp-hosted/src/proto.rs
index 74c67bd61..09bec8984 100644
--- a/embassy-net-esp-hosted/src/proto.rs
+++ b/embassy-net-esp-hosted/src/proto.rs
@@ -16,7 +16,7 @@ Switch to a proper script when https://github.com/YuhanLiin/micropb/issues/30 is
     // Special config for things that need to be larger
     g.configure(
         ".CtrlMsg_Req_OTAWrite.ota_data",
-        micropb_gen::Config::new().max_bytes(1024),
+        micropb_gen::Config::new().max_bytes(256),
     );
     g.configure(
         ".CtrlMsg_Event_ESPInit.init_data",
@@ -4296,28 +4296,28 @@ impl ::micropb::MessageEncode for CtrlMsg_Resp_OTABegin {
 #[derive(Debug, Default, PartialEq, Clone)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub struct CtrlMsg_Req_OTAWrite {
-    pub r#ota_data: ::micropb::heapless::Vec<u8, 1024>,
+    pub r#ota_data: ::micropb::heapless::Vec<u8, 256>,
 }
 impl CtrlMsg_Req_OTAWrite {
     ///Return a reference to `ota_data`
     #[inline]
-    pub fn r#ota_data(&self) -> &::micropb::heapless::Vec<u8, 1024> {
+    pub fn r#ota_data(&self) -> &::micropb::heapless::Vec<u8, 256> {
         &self.r#ota_data
     }
     ///Return a mutable reference to `ota_data`
     #[inline]
-    pub fn mut_ota_data(&mut self) -> &mut ::micropb::heapless::Vec<u8, 1024> {
+    pub fn mut_ota_data(&mut self) -> &mut ::micropb::heapless::Vec<u8, 256> {
         &mut self.r#ota_data
     }
     ///Set the value of `ota_data`
     #[inline]
-    pub fn set_ota_data(&mut self, value: ::micropb::heapless::Vec<u8, 1024>) -> &mut Self {
+    pub fn set_ota_data(&mut self, value: ::micropb::heapless::Vec<u8, 256>) -> &mut Self {
         self.r#ota_data = value.into();
         self
     }
     ///Builder method that sets the value of `ota_data`. Useful for initializing the message.
     #[inline]
-    pub fn init_ota_data(mut self, value: ::micropb::heapless::Vec<u8, 1024>) -> Self {
+    pub fn init_ota_data(mut self, value: ::micropb::heapless::Vec<u8, 256>) -> Self {
         self.r#ota_data = value.into();
         self
     }