fix: answer requests the state machine was about to drop

updatehub/src/cloud_mock.rs

@@ -13,6 +13,12 @@ std::thread_local! {
     static OBJECT_DATA: RefCell<Option<Vec<u8>>> = const { RefCell::new(Option::None) };
 }
 
+type ProbeGate = (async_channel::Receiver<()>, async_channel::Sender<()>);
+
+std::thread_local! {
+    static PROBE_GATE: RefCell<Option<ProbeGate>> = const { RefCell::new(None) };
+}
+
 pub(crate) enum FakeResponse {
     NoUpdate,
     HasUpdate,
@@ -28,6 +34,18 @@ pub(crate) fn setup_fake_response(res: FakeResponse) {
     RESPONSE_CONFIG.with(|conf| conf.replace_with(move |&mut _| res));
 }
 
+/// Holds every probe until it is released, standing in for a server that has
+/// not answered yet.
+///
+/// Returns the release side and a receiver that reports each probe reaching the
+/// server, so a test can act while a request is provably half-answered.
+pub(crate) fn gate_probes() -> (async_channel::Sender<()>, async_channel::Receiver<()>) {
+    let (release, gate) = async_channel::unbounded();
+    let (reached, arrived) = async_channel::unbounded();
+    PROBE_GATE.with(|conf| conf.borrow_mut().replace((gate, reached)));
+    (release, arrived)
+}
+
 pub(crate) fn set_download_data(data: Vec<u8>) {
     OBJECT_DATA.with(|conf| conf.borrow_mut().replace(data));
 }
@@ -42,6 +60,14 @@ impl<'a> Client<'a> {
         _num_retries: usize,
         _firmware: api::FirmwareMetadata<'_>,
     ) -> Result<api::ProbeResponse> {
+        // Cloned out before awaiting, so the `RefCell` is never borrowed across
+        // a suspension point.
+        let gate = PROBE_GATE.with(|conf| conf.borrow().clone());
+        if let Some((gate, reached)) = gate {
+            let _ = reached.send(()).await;
+            let _ = gate.recv().await;
+        }
+
         RESPONSE_CONFIG.with(|conf| match std::ops::Deref::deref(&conf.borrow()) {
             FakeResponse::NoUpdate => Ok(api::ProbeResponse::NoUpdate),
             FakeResponse::ExtraPoll => Ok(api::ProbeResponse::ExtraPoll(10)),

updatehub/src/states/direct_download.rs

@@ -29,10 +29,9 @@ impl StateChangeImpl for DirectDownload {
     async fn handle(self, context: &mut Context) -> Result<(State, machine::StepTransition)> {
         info!("fetching update package directly from url: {:?}", self.url);
 
-        let communication_receiver = &context.communication.receiver.clone();
         let context = Mutex::new(context);
 
-        let download_future = async {
+        let download = async {
             let download_dir = context.lock().await.settings.update.download_dir.clone();
             tokio::fs::create_dir_all(&download_dir)
                 .await
@@ -46,27 +45,8 @@ impl StateChangeImpl for DirectDownload {
             Ok(State::PrepareLocalInstall(PrepareLocalInstall { update_file }))
         };
 
-        let message_handle_future = async {
-            while let Ok((msg, responder)) = communication_receiver.recv().await {
-                if let Some(new_state) =
-                    self.handle_communication(msg, responder, *context.lock().await).await
-                {
-                    return Ok(new_state);
-                }
-            }
+        let state = self.handle_communication_while(&context, download).await?;
 
-            Err(super::TransitionError::CommunicationFailed)
-        };
-
-        futures_util::pin_mut!(download_future);
-        futures_util::pin_mut!(message_handle_future);
-
-        Ok((
-            futures_util::future::select(download_future, message_handle_future)
-                .await
-                .factor_first()
-                .0?,
-            machine::StepTransition::Immediate,
-        ))
+        Ok((state, machine::StepTransition::Immediate))
     }
 }

updatehub/src/states/download.rs

@@ -23,7 +23,7 @@ pub(super) struct Download {
 
 impl Download {
     async fn start_download(
-        update_package: UpdatePackage,
+        update_package: &UpdatePackage,
         context: &Mutex<&mut Context>,
     ) -> Result<()> {
         let installation_set =
@@ -129,52 +129,20 @@ impl StateChangeImpl for Download {
     }
 
     async fn handle(self, context: &mut Context) -> Result<(State, machine::StepTransition)> {
-        let communication_receiver = &context.communication.receiver.clone();
         let context = Mutex::new(context);
 
-        let update_package = self.update_package.clone();
-        let download_future = async {
-            Download::start_download(update_package.clone(), &context).await?;
-            Result::Ok(None)
+        let download = async {
+            Download::start_download(&self.update_package, &context).await?;
+            Ok(State::Validation(Validation {
+                package: self.update_package.clone(),
+                sign: self.sign.clone(),
+                require_download: false,
+            }))
         };
 
-        let message_handle_future = async {
-            while let Ok((msg, responder)) = communication_receiver.recv().await {
-                if let Some(new_state) =
-                    self.handle_communication(msg, responder, *context.lock().await).await
-                {
-                    return Ok(Some(new_state));
-                }
-            }
-            Ok(None)
-        };
+        let state = self.handle_communication_while(&context, download).await?;
 
-        // Clone update package and object_context so self can be freely held by
-        // message_handle_future
-        let update_package = self.update_package.clone();
-        let sign = self.sign.clone();
-
-        // download_future dones't need to be pinned as it doesn't borrow context
-        futures_util::pin_mut!(download_future);
-        futures_util::pin_mut!(message_handle_future);
-
-        if let Some(new_state) =
-            futures_util::future::select(download_future, message_handle_future)
-                .await
-                .factor_first()
-                .0?
-        {
-            return Ok((new_state, machine::StepTransition::Immediate));
-        }
-
-        Ok((
-            State::Validation(Validation {
-                package: update_package,
-                sign,
-                require_download: false,
-            }),
-            machine::StepTransition::Immediate,
-        ))
+        Ok((state, machine::StepTransition::Immediate))
     }
 }
 

updatehub/src/states/machine/address.rs

@@ -49,43 +49,50 @@ pub(crate) enum StateResponse {
     InvalidState(String),
 }
 
+// Either half of a request's channel going away means the state machine is no
+// longer running. Report that rather than taking the caller's task down with a
+// panic.
 impl<T> From<async_channel::SendError<T>> for crate::states::TransitionError {
-    fn from(err: async_channel::SendError<T>) -> Self {
-        unreachable!("Unexpected sending error for {:?}", err)
+    fn from(_: async_channel::SendError<T>) -> Self {
+        crate::states::TransitionError::CommunicationFailed
+    }
+}
+
+impl From<async_channel::RecvError> for crate::states::TransitionError {
+    fn from(_: async_channel::RecvError) -> Self {
+        crate::states::TransitionError::CommunicationFailed
     }
 }
 
 impl Addr {
-    pub(crate) async fn request_info(&self) -> super::Result<sdk::api::info::Response> {
+    /// Hands `msg` to the state machine and waits for the answer to come back.
+    async fn request(&self, msg: Message) -> super::Result<Response> {
         let (sndr, recv) = async_channel::bounded(1);
-        self.message.send((Message::Info, sndr)).await?;
-        match recv.recv().await {
-            Ok(Ok(Response::Info(resp))) => Ok(*resp),
-            Ok(Err(e)) => Err(e),
-            res => unreachable!("Unexpected response: {:?}", res),
+        self.message.send((msg, sndr)).await?;
+        recv.recv().await?
+    }
+
+    pub(crate) async fn request_info(&self) -> super::Result<sdk::api::info::Response> {
+        match self.request(Message::Info).await? {
+            Response::Info(resp) => Ok(*resp),
+            res => unreachable!("Unexpected response: {res:?}"),
         }
     }
 
     pub(crate) async fn request_probe(
         &self,
         custom_server: Option<String>,
     ) -> super::Result<ProbeResponse> {
-        let (sndr, recv) = async_channel::bounded(1);
-        self.message.send((Message::Probe(custom_server), sndr)).await?;
-        match recv.recv().await {
-            Ok(Ok(Response::Probe(resp))) => Ok(resp),
-            Ok(Err(e)) => Err(e),
-            res => unreachable!("Unexpected response: {:?}", res),
+        match self.request(Message::Probe(custom_server)).await? {
+            Response::Probe(resp) => Ok(resp),
+            res => unreachable!("Unexpected response: {res:?}"),
         }
     }
 
     pub(crate) async fn request_abort_download(&self) -> super::Result<AbortDownloadResponse> {
-        let (sndr, recv) = async_channel::bounded(1);
-        self.message.send((Message::AbortDownload, sndr)).await?;
-        match recv.recv().await {
-            Ok(Ok(Response::AbortDownload(resp))) => Ok(resp),
-            Ok(Err(e)) => Err(e),
-            res => unreachable!("Unexpected response: {:?}", res),
+        match self.request(Message::AbortDownload).await? {
+            Response::AbortDownload(resp) => Ok(resp),
+            res => unreachable!("Unexpected response: {res:?}"),
         }
     }
 
@@ -94,23 +101,17 @@ impl Addr {
         path: PathBuf,
     ) -> super::Result<StateResponse> {
         trace!("Local install requested");
-        let (sndr, recv) = async_channel::bounded(1);
-        self.message.send((Message::LocalInstall(path), sndr)).await?;
-        match recv.recv().await {
-            Ok(Ok(Response::LocalInstall(resp))) => Ok(resp),
-            Ok(Err(e)) => Err(e),
-            res => unreachable!("Unexpected response: {:?}", res),
+        match self.request(Message::LocalInstall(path)).await? {
+            Response::LocalInstall(resp) => Ok(resp),
+            res => unreachable!("Unexpected response: {res:?}"),
         }
     }
 
     pub(crate) async fn request_remote_install(&self, url: String) -> super::Result<StateResponse> {
         trace!("Remote install requested");
-        let (sndr, recv) = async_channel::bounded(1);
-        self.message.send((Message::RemoteInstall(url), sndr)).await?;
-        match recv.recv().await {
-            Ok(Ok(Response::RemoteInstall(resp))) => Ok(resp),
-            Ok(Err(e)) => Err(e),
-            res => unreachable!("Unexpected response: {:?}", res),
+        match self.request(Message::RemoteInstall(url)).await? {
+            Response::RemoteInstall(resp) => Ok(resp),
+            res => unreachable!("Unexpected response: {res:?}"),
         }
     }
 }