Support waiting on a call with optional response (#611)

Part of #577. This is going to be used by `AppletExitStatus` too.

crates/cli-tools/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Minor
 
+- Change the flags of `action::AppletRpc` to use `action::Wait`
+- Add `action::Wait` for commands returning an optional response
 - Increase the default connection timeout from 1 to 5 seconds
 - Add `action::PlatformUpdate` for platform update
 - Add `action::PlatformList` to list connected platforms
@@ -20,4 +22,4 @@
 
 ## 0.1.0
 
-<!-- Increment to skip CHANGELOG.md test: 4 -->
+<!-- Increment to skip CHANGELOG.md test: 5 -->

crates/cli-tools/Cargo.toml

@@ -21,13 +21,13 @@ indicatif = { version = "0.17.8", default-features = false }
 rusb = { version = "0.9.4", default-features = false }
 serde = { version = "1.0.202", default-features = false, features = ["derive"] }
 toml = { version = "0.8.13", default-features = false, features = ["display", "parse"] }
-wasefire-protocol = { version = "0.1.1-git", path = "../protocol", features = ["host"] }
+wasefire-protocol = { version = "0.2.0-git", path = "../protocol", features = ["host"] }
 wasefire-wire = { version = "0.1.1-git", path = "../wire" }
 
 [dependencies.tokio]
 version = "1.40.0"
 default-features = false
-features = ["fs", "io-std", "process", "rt"]
+features = ["fs", "io-std", "process", "rt", "time"]
 
 [dependencies.wasefire-protocol-tokio]
 version = "0.1.0-git"

crates/cli-tools/src/action.rs

@@ -14,14 +14,15 @@
 
 use std::fmt::Display;
 use std::path::{Path, PathBuf};
+use std::time::Duration;
 
 use anyhow::{bail, ensure, Result};
 use cargo_metadata::{Metadata, MetadataCommand};
 use clap::{ValueEnum, ValueHint};
 use rusb::GlobalContext;
 use tokio::process::Command;
 use wasefire_protocol::{self as service, applet, Connection, ConnectionExt};
-use wasefire_wire::Yoke;
+use wasefire_wire::{self as wire, Yoke};
 
 use crate::{cmd, fs};
 
@@ -90,27 +91,72 @@ pub struct AppletRpc {
     #[clap(flatten)]
     rpc: Rpc,
 
-    /// Number of retries to receive a response.
-    #[arg(long, default_value = "3")]
-    retries: usize,
+    #[clap(flatten)]
+    wait: Wait,
 }
 
 impl AppletRpc {
     pub async fn run(self, connection: &mut dyn Connection) -> Result<()> {
-        let AppletRpc { applet, rpc, retries } = self;
+        let AppletRpc { applet, rpc, mut wait } = self;
         let applet_id = match applet {
             Some(_) => bail!("applet identifiers are not supported yet"),
             None => applet::AppletId,
         };
         let request = applet::Request { applet_id, request: &rpc.read().await? };
         connection.call::<service::AppletRequest>(request).await?.get();
-        for _ in 0 .. retries {
-            let response = connection.call::<service::AppletResponse>(applet_id).await?;
-            if let Some(response) = response.get().response {
-                return rpc.write(response).await;
+        wait.ensure_wait();
+        match wait.run::<service::AppletResponse, &[u8]>(connection, applet_id).await? {
+            None => bail!("did not receive a response"),
+            Some(response) => rpc.write(response.get()).await,
+        }
+    }
+}
+
+/// Options to repeatedly call a command with an optional response.
+#[derive(clap::Args)]
+pub struct Wait {
+    /// Waits until there is a response.
+    ///
+    /// This is equivalent to --period=100ms.
+    #[arg(long)]
+    wait: bool,
+
+    /// Retries every so often until there is a response.
+    ///
+    /// The command doesn't return `None` in that case.
+    #[arg(long, conflicts_with = "wait")]
+    period: Option<humantime::Duration>,
+}
+
+impl Wait {
+    pub fn ensure_wait(&mut self) {
+        if self.wait || self.period.is_some() {
+            return;
+        }
+        self.wait = true;
+    }
+
+    pub async fn run<S, T: wire::Wire<'static>>(
+        self, connection: &mut dyn Connection, request: S::Request<'_>,
+    ) -> Result<Option<Yoke<T::Type<'static>>>>
+    where S: for<'a> service::Service<Response<'a> = Option<T::Type<'a>>> {
+        let Wait { wait, period } = self;
+        let period = match (wait, period) {
+            (true, None) => Some(Duration::from_millis(100)),
+            (true, Some(_)) => unreachable!(),
+            (false, None) => None,
+            (false, Some(x)) => Some(*x),
+        };
+        let request = S::request(request);
+        loop {
+            match connection.call_ref::<S>(&request).await?.try_map(|x| x.ok_or(())) {
+                Ok(x) => break Ok(Some(x)),
+                Err(()) => match period {
+                    Some(period) => tokio::time::sleep(period).await,
+                    None => break Ok(None),
+                },
             }
         }
-        bail!("did not receive a response after {retries} retries");
     }
 }
 

crates/protocol-tokio/CHANGELOG.md

@@ -2,4 +2,4 @@
 
 ## 0.1.0-git
 
-<!-- Increment to skip CHANGELOG.md test: 3 -->
+<!-- Increment to skip CHANGELOG.md test: 4 -->

crates/protocol-tokio/Cargo.toml

@@ -31,7 +31,7 @@ features = ["std"]
 optional = true
 
 [dependencies.wasefire-protocol]
-version = "0.1.1-git"
+version = "0.2.0-git"
 path = "../protocol"
 features = ["host"]
 optional = true

crates/protocol-usb/CHANGELOG.md

@@ -17,4 +17,4 @@
 
 ## 0.1.0
 
-<!-- Increment to skip CHANGELOG.md test: 3 -->
+<!-- Increment to skip CHANGELOG.md test: 4 -->

crates/protocol-usb/Cargo.toml

@@ -25,7 +25,7 @@ path = "../board"
 optional = true
 
 [dependencies.wasefire-protocol]
-version = "0.1.1-git"
+version = "0.2.0-git"
 path = "../protocol"
 features = ["host"]
 optional = true

crates/protocol/CHANGELOG.md

@@ -1,9 +1,14 @@
 # Changelog
 
-## 0.1.1-git
+## 0.2.0-git
+
+### Major
+
+- Remove `applet::Response` and inline its definition in `AppletResponse::Response`
 
 ### Minor
 
+- Add `ConnectionExt::call_ref()` to share a request between calls
 - Add a `Service::NAME` constant with `host` feature
 - Add `PlatformUpdate{Metadata,Transfer}` calls and `transfer` module for platform updates
 - Add a `Connection` abstraction with `host` feature

crates/protocol/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wasefire-protocol"
-version = "0.1.1-git"
+version = "0.2.0-git"
 authors = ["Julien Cretin <cretin@google.com>"]
 license = "Apache-2.0"
 publish = true

crates/protocol/crates/schema/device.txt

@@ -2,7 +2,7 @@ result: {Ok=0:() Err=1:(space:u8 code:u16)}
 version: 3
 0 [0 -] ApiVersion: () -> u32
 1 [0 -] AppletRequest: (applet_id:() request:[u8]) -> ()
-2 [0 -] AppletResponse: () -> (response:{None=0:() Some=1:[u8]})
+2 [0 -] AppletResponse: () -> {None=0:() Some=1:[u8]}
 3 [0 -] PlatformReboot: () -> {}
 4 [0 -] AppletTunnel: (applet_id:() delimiter:[u8]) -> ()
 5 [1 -] PlatformInfo: () -> (serial:[u8] version:[u8])

crates/protocol/crates/schema/host.txt

@@ -2,7 +2,7 @@ result: {Ok=0:() Err=1:(space:u8 code:u16)}
 version: 3
 0 [0 -] ApiVersion: () -> u32
 1 [0 -] AppletRequest: (applet_id:() request:[u8]) -> ()
-2 [0 -] AppletResponse: () -> (response:{None=0:() Some=1:[u8]})
+2 [0 -] AppletResponse: () -> {None=0:() Some=1:[u8]}
 3 [0 -] PlatformReboot: () -> {}
 4 [0 -] AppletTunnel: (applet_id:() delimiter:[u8]) -> ()
 5 [1 -] PlatformInfo: () -> (serial:[u8] version:[u8])

crates/protocol/src/applet.rs

@@ -20,11 +20,6 @@ pub struct Request<'a> {
     pub request: &'a [u8],
 }
 
-#[derive(Debug, Wire)]
-pub struct Response<'a> {
-    pub response: Option<&'a [u8]>,
-}
-
 #[derive(Debug, Copy, Clone, Wire)]
 pub struct AppletId;
 

crates/protocol/src/connection.rs

@@ -46,11 +46,15 @@ pub trait ConnectionExt: Connection {
     /// Calls a service on the device.
     fn call<S: Service>(
         &mut self, request: S::Request<'_>,
+    ) -> impl Future<Output = anyhow::Result<Yoke<S::Response<'static>>>> {
+        async { self.call_ref::<S>(&S::request(request)).await }
+    }
+
+    fn call_ref<S: Service>(
+        &mut self, request: &Api<Request>,
     ) -> impl Future<Output = anyhow::Result<Yoke<S::Response<'static>>>> {
         async {
-            self.send(&S::request(request))
-                .await
-                .with_context(|| format!("sending {}", S::NAME))?;
+            self.send(request).await.with_context(|| format!("sending {}", S::NAME))?;
             self.receive::<S>().await.with_context(|| format!("receiving {}", S::NAME))
         }
     }

crates/protocol/src/lib.rs

@@ -220,7 +220,7 @@ api! {
     1 [0 -] AppletRequest: applet::Request<'a> => (),
 
     /// Reads a response from an applet.
-    2 [0 -] AppletResponse: applet::AppletId => applet::Response<'a>,
+    2 [0 -] AppletResponse: applet::AppletId => Option<&'a [u8]>,
 
     /// Reboots the platform.
     3 [0 -] PlatformReboot: () => !,

crates/scheduler/CHANGELOG.md

@@ -160,4 +160,4 @@
 
 ## 0.1.0
 
-<!-- Increment to skip CHANGELOG.md test: 4 -->
+<!-- Increment to skip CHANGELOG.md test: 5 -->

crates/scheduler/Cargo.toml

@@ -36,7 +36,7 @@ features = ["toctou"]
 optional = true
 
 [dependencies.wasefire-protocol]
-version = "0.1.1-git"
+version = "0.2.0-git"
 path = "../protocol"
 features = ["device"]
 

crates/scheduler/src/protocol.rs

@@ -75,8 +75,7 @@ fn process_event_<B: Board>(
         }
         Api::AppletResponse(applet_id) => {
             let response = applet::<B>(scheduler, applet_id)?.get_response()?;
-            let response = response.as_deref();
-            reply::<B, service::AppletResponse>(service::applet::Response { response });
+            reply::<B, service::AppletResponse>(response.as_deref());
         }
         Api::PlatformReboot(()) => {
             use wasefire_board_api::platform::Api as _;

examples/rust/protocol/host/src/main.rs

@@ -56,7 +56,7 @@ async fn main() -> Result<()> {
             connection.call::<service::AppletRequest>(request).await?.get();
             loop {
                 let response = connection.call::<service::AppletResponse>(AppletId).await?;
-                if let Some(response) = response.get().response {
+                if let Some(response) = response.get() {
                     print!("{}", std::str::from_utf8(response).unwrap());
                     break Ok(());
                 }