From a81c17cbb1bf0cd6a03a98edd6f0a2d05b1300fd Mon Sep 17 00:00:00 2001
From: Dreaded_X <tim@huizinga.dev>
Date: Sat, 16 Sep 2023 03:58:06 +0200
Subject: [PATCH] Small tweaks

---
 updater/Cargo.lock |  4 ++++
 updater/Cargo.toml |  2 ++
 updater/src/lib.rs | 21 ++++++++++++---------
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/updater/Cargo.lock b/updater/Cargo.lock
index d30ae52..d8d72fa 100644
--- a/updater/Cargo.lock
+++ b/updater/Cargo.lock
@@ -378,6 +378,9 @@ dependencies = [
 name = "embassy-futures"
 version = "0.1.0"
 source = "git+https://github.com/embassy-rs/embassy#f1f4943ca51e8827146daca950fdf88d5b1e046b"
+dependencies = [
+ "defmt",
+]
 
 [[package]]
 name = "embassy-net"
@@ -1215,6 +1218,7 @@ dependencies = [
  "defmt",
  "defmt-rtt",
  "embassy-boot",
+ "embassy-futures",
  "embassy-net",
  "embassy-time",
  "embedded-io-async",
diff --git a/updater/Cargo.toml b/updater/Cargo.toml
index dd4848f..aca1b53 100644
--- a/updater/Cargo.toml
+++ b/updater/Cargo.toml
@@ -27,6 +27,7 @@ embassy-time = { version = "0.1", features = [
 	"defmt-timestamp-uptime",
 	"nightly",
 ] }
+embassy-futures = { version = "0.1", features = ["defmt"] }
 rand_core = "0.6.4"
 embedded-io-async = { version = "0.5", features = ["defmt-03"] }
 embedded-storage = "0.3.0"
@@ -48,6 +49,7 @@ static_cell = { version = "1.2.0", features = ["nightly"] }
 impl-tools = "0.10.0"
 
 [patch.crates-io]
+embassy-futures = { git = "https://github.com/embassy-rs/embassy" }
 embassy-net = { git = "https://github.com/embassy-rs/embassy" }
 embassy-boot = { git = "https://github.com/embassy-rs/embassy" }
 embassy-time = { git = "https://github.com/embassy-rs/embassy" }
diff --git a/updater/src/lib.rs b/updater/src/lib.rs
index 706e82b..c956060 100644
--- a/updater/src/lib.rs
+++ b/updater/src/lib.rs
@@ -30,7 +30,7 @@ pub use crate::error::Error;
 
 #[derive(Serialize)]
 #[serde(rename_all = "snake_case", tag = "status")]
-enum Status<'a> {
+pub enum Status<'a> {
     Connected { version: &'a str },
     Disconnected,
     PreparingUpdate,
@@ -48,6 +48,9 @@ impl Status<'_> {
     }
 }
 
+/// This is a wrapper around `BlockingFirmwareUpdater` that downloads signed updates
+/// from a HTTPS url.
+/// It also provides the current device status over MQTT
 // TODO: Make this the owner of the blocking firmware updater
 // TODO: When fixed, use the async firmware updater
 pub struct Updater<'a, DFU, STATE>
@@ -59,7 +62,6 @@ where
     updater: BlockingFirmwareUpdater<'a, DFU, STATE>,
 
     topic_status: &'static str,
-    topic_update: &'static str,
     version: &'static str,
     public_key: &'static [u8],
 }
@@ -70,22 +72,22 @@ where
     STATE: NorFlash,
     DFU::Error: Format,
 {
+    /// Wrap the `BlockingFirmwareUpdater`
     pub fn new(
         updater: BlockingFirmwareUpdater<'a, DFU, STATE>,
         topic_status: &'static str,
-        topic_update: &'static str,
         version: &'static str,
         public_key: &'static [u8],
     ) -> Self {
         Self {
             updater,
             topic_status,
-            topic_update,
             version,
             public_key,
         }
     }
 
+    /// Set MQTT connection up to notify over MQTT when the device loses connection
     pub fn add_will<const MAX_PROPERTIES: usize>(
         &self,
         config: &mut ClientConfig<'_, MAX_PROPERTIES, impl RngCore>,
@@ -94,6 +96,8 @@ where
         config.add_will(self.topic_status, msg, true);
     }
 
+    /// Mark the device is ready and booted, will notify over MQTT that the device is connected and the
+    /// currently running firmware version
     pub async fn ready<const MAX_PROPERTIES: usize>(
         &mut self,
         client: &mut MqttClient<'_, impl Write + Read, MAX_PROPERTIES, impl RngCore>,
@@ -107,21 +111,20 @@ where
             .send_message(self.topic_status, &status, QualityOfService::QoS1, true)
             .await?;
 
-        client.subscribe_to_topic(self.topic_update).await?;
-
         self.updater.mark_booted()?;
 
         Ok(())
     }
 
+    /// Download signed update from specified url and notify progress over MQTT
     pub async fn update<const MAX_PROPERTIES: usize>(
         &mut self,
+        url: Url<'_>,
         stack: &'static Stack<impl Driver>,
         rng: &mut (impl RngCore + CryptoRng),
         client: &mut MqttClient<'_, impl Write + Read, MAX_PROPERTIES, impl RngCore>,
-        url: Url<'_>,
     ) -> Result<!, Error<DFU::Error>> {
-        let result = self._update(stack, rng, client, url).await;
+        let result = self._update(url, stack, rng, client).await;
 
         if let Err(err) = &result {
             let status = Status::UpdateFailed {
@@ -139,10 +142,10 @@ where
 
     async fn _update<const MAX_PROPERTIES: usize>(
         &mut self,
+        url: Url<'_>,
         stack: &'static Stack<impl Driver>,
         rng: &mut (impl RngCore + CryptoRng),
         client: &mut MqttClient<'_, impl Write + Read, MAX_PROPERTIES, impl RngCore>,
-        url: Url<'_>,
     ) -> Result<!, Error<DFU::Error>> {
         info!("Preparing for OTA...");
         let status = Status::PreparingUpdate.json();