From a249859e7f5d3d2c04cf467b16fea133c5f8427d Mon Sep 17 00:00:00 2001
From: ivmarkov <ivan.markov@gmail.com>
Date: Thu, 21 Sep 2023 18:34:28 +0000
Subject: [PATCH] Option to postpone setting the boot partition after update

---
 src/ota.rs | 30 +++++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/ota.rs b/src/ota.rs
index 8672d2e..bcd9823 100644
--- a/src/ota.rs
+++ b/src/ota.rs
@@ -141,12 +141,16 @@ where
 }
 
 pub trait OtaUpdate: Write {
-    fn complete(&mut self) -> Result<(), Self::Error>;
+    type OtaUpdateFinished: OtaUpdateFinished;
 
-    fn abort(&mut self) -> Result<(), Self::Error>;
+    fn finish(self) -> Result<Self::OtaUpdateFinished, Self::Error>;
+
+    fn complete(self) -> Result<(), Self::Error>;
+
+    fn abort(self) -> Result<(), Self::Error>;
 
     fn update<R>(
-        &mut self,
+        mut self,
         read: R,
         progress: impl Fn(u64, u64),
     ) -> Result<(), CopyError<R::Error, Self::Error>>
@@ -156,7 +160,7 @@ pub trait OtaUpdate: Write {
     {
         let mut buf = [0_u8; 64];
 
-        match copy_len_with_progress(read, &mut *self, &mut buf, u64::MAX, progress) {
+        match copy_len_with_progress(read, &mut self, &mut buf, u64::MAX, progress) {
             Ok(_) => self.complete().map_err(CopyError::Write),
             Err(e) => {
                 self.abort().map_err(CopyError::Write)?;
@@ -167,6 +171,10 @@ pub trait OtaUpdate: Write {
     }
 }
 
+pub trait OtaUpdateFinished: ErrorType {
+    fn activate(self) -> Result<(), Self::Error>;
+}
+
 #[cfg(feature = "nightly")]
 pub mod asynch {
     use crate::io::asynch::{ErrorType, Read, Write};
@@ -236,12 +244,16 @@ pub mod asynch {
     }
 
     pub trait OtaUpdate: Write {
-        async fn complete(&mut self) -> Result<(), Self::Error>;
+        type OtaUpdateFinished: OtaUpdateFinished;
 
-        async fn abort(&mut self) -> Result<(), Self::Error>;
+        async fn finish(self) -> Result<Self::OtaUpdateFinished, Self::Error>;
+
+        async fn complete(self) -> Result<(), Self::Error>;
+
+        async fn abort(self) -> Result<(), Self::Error>;
 
         async fn update<R>(
-            &mut self,
+            self,
             read: R,
             progress: impl Fn(u64, u64),
         ) -> Result<(), CopyError<R::Error, Self::Error>>
@@ -249,4 +261,8 @@ pub mod asynch {
             R: Read,
             Self: Sized;
     }
+
+    pub trait OtaUpdateFinished: ErrorType {
+        async fn activate(self) -> Result<(), Self::Error>;
+    }
 }