Skip to main content

wasefire_protocol/
connection.rs

1// Copyright 2024 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use alloc::boxed::Box;
16use alloc::format;
17use core::any::Any;
18use core::future::Future;
19use core::ops::Deref;
20use core::pin::Pin;
21
22use anyhow::{Context, ensure};
23use wasefire_wire::Yoke;
24
25use crate::{Api, ApiResult, Request, Service, VERSION};
26
27/// Connection with a cached API version.
28#[derive(PartialEq, Eq)]
29pub struct Device<T: Connection> {
30    version: u32,
31    connection: T,
32}
33
34pub type DynDevice = Device<Box<dyn Connection>>;
35
36impl<T: Connection> Device<T> {
37    /// Checks whether the device is supported and caches the API version.
38    pub async fn new(connection: T) -> anyhow::Result<Self> {
39        let version = *connection.call::<crate::ApiVersion>(()).await?.get();
40        ensure!(version <= VERSION, "the device is more recent than the host");
41        Ok(Device { version, connection })
42    }
43
44    /// Returns the cached API version.
45    pub fn version(&self) -> u32 {
46        self.version
47    }
48
49    /// Extracts the connection.
50    pub fn connection(self) -> T {
51        self.connection
52    }
53
54    /// Returns whether the device supports a service.
55    pub fn supports<S: Service>(&self) -> bool {
56        // The device is not newer than the host by invariant.
57        S::VERSIONS.contains(self.version).unwrap()
58    }
59
60    pub async fn platform_info(&self) -> anyhow::Result<crate::platform::DynInfo> {
61        if self.supports::<crate::PlatformInfo3>() {
62            Ok(crate::platform::DynInfo::V3(self.call::<crate::PlatformInfo3>(()).await?))
63        } else if self.supports::<crate::_PlatformInfo2>() {
64            Ok(crate::platform::DynInfo::V2(self.call::<crate::_PlatformInfo2>(()).await?))
65        } else if self.supports::<crate::_PlatformInfo1>() {
66            Ok(crate::platform::DynInfo::V1(self.call::<crate::_PlatformInfo1>(()).await?))
67        } else if self.supports::<crate::_PlatformInfo0>() {
68            Ok(crate::platform::DynInfo::V0(self.call::<crate::_PlatformInfo0>(()).await?))
69        } else {
70            Err(wasefire_error::Error::world(wasefire_error::Code::NotImplemented).into())
71        }
72    }
73}
74
75impl<T: Connection> Deref for Device<T> {
76    type Target = T;
77
78    fn deref(&self) -> &Self::Target {
79        &self.connection
80    }
81}
82
83pub type DynFuture<'a, T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + 'a>>;
84
85pub trait Connection: Any {
86    /// Sends a raw request (possibly tunneled) to the device.
87    fn write<'a>(&'a self, request: &'a [u8]) -> DynFuture<'a, ()>;
88
89    /// Receives a raw response (possibly tunneled) from the device.
90    fn read(&self) -> DynFuture<'_, Box<[u8]>>;
91}
92
93impl Connection for Box<dyn Connection> {
94    fn write<'a>(&'a self, request: &'a [u8]) -> DynFuture<'a, ()> {
95        (**self).write(request)
96    }
97
98    fn read(&self) -> DynFuture<'_, Box<[u8]>> {
99        (**self).read()
100    }
101}
102
103pub trait ConnectionExt: Connection {
104    /// Calls a service on the device.
105    fn call<S: Service>(
106        &self, request: S::Request<'_>,
107    ) -> impl Future<Output = anyhow::Result<Yoke<S::Response<'static>>>> {
108        async { self.call_ref::<S>(&S::request(request)).await }
109    }
110
111    fn call_ref<S: Service>(
112        &self, request: &Api<Request>,
113    ) -> impl Future<Output = anyhow::Result<Yoke<S::Response<'static>>>> {
114        async {
115            self.send(request).await.with_context(|| format!("sending {}", S::NAME))?;
116            self.receive::<S>().await.with_context(|| format!("receiving {}", S::NAME))
117        }
118    }
119
120    /// Sends a request to the device.
121    fn send(&self, request: &Api<'_, Request>) -> impl Future<Output = anyhow::Result<()>> {
122        async {
123            let request = request.encode().context("encoding request")?;
124            self.write(&request).await
125        }
126    }
127
128    /// Receives a response from the device.
129    fn receive<S: Service>(
130        &self,
131    ) -> impl Future<Output = anyhow::Result<Yoke<S::Response<'static>>>> {
132        async {
133            let response = self.read().await?;
134            let response = ApiResult::<S>::decode_yoke(response).context("decoding response")?;
135            response.try_map(|x| match x {
136                ApiResult::Ok(x) => Ok(x),
137                ApiResult::Err(error) => Err(anyhow::Error::new(error)),
138            })
139        }
140    }
141}
142
143impl<T: Connection + ?Sized> ConnectionExt for T {}