Skip to content

Commit

Permalink
virt: Add possiblity to run with multiple clients
Browse files Browse the repository at this point in the history
  • Loading branch information
sosthene-nitrokey committed Nov 9, 2023
1 parent 00e4d21 commit 9763723
Showing 1 changed file with 85 additions and 8 deletions.
93 changes: 85 additions & 8 deletions src/virt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,18 @@
mod store;
mod ui;

use std::{path::PathBuf, sync::Mutex};
use std::{
path::PathBuf,
sync::{Arc, Mutex},
};

use rand_chacha::ChaCha8Rng;
use rand_core::SeedableRng as _;

use crate::{
backend::{BackendId, CoreOnly, Dispatch},
client::ClientBuilder,
platform,
platform::{self, Syscall},
service::Service,
ClientImplementation,
};
Expand All @@ -23,6 +26,9 @@ pub use store::{Filesystem, Ram, StoreProvider};
pub use ui::UserInterface;

pub type Client<S, D = CoreOnly> = ClientImplementation<Service<Platform<S>, D>, D>;
/// Virtual client that can coexist with other clients
pub type MultiClient<S, D = CoreOnly> =
ClientImplementation<Arc<Mutex<Service<Platform<S>, D>>>, D>;

// We need this mutex to make sure that:
// - TrussedInterchange is not used concurrently (panics if violated)
Expand Down Expand Up @@ -71,18 +77,59 @@ where
with_client(Ram::default(), client_id, f)
}

pub fn with_clients<S, R, F, const N: usize>(store: S, client_ids: [&str; N], f: F) -> R
where
S: StoreProvider,
F: FnOnce([MultiClient<S>; N]) -> R,
{
with_platform(store, |platform| platform.run_clients(client_ids, f))
}

pub fn with_fs_clients<P, R, F, const N: usize>(internal: P, client_ids: [&str; N], f: F) -> R
where
P: Into<PathBuf>,
F: FnOnce([MultiClient<Filesystem>; N]) -> R,
{
with_clients(Filesystem::new(internal), client_ids, f)
}

/// Run a function with multiple clients using the RAM for the filesystem.
///
///
/// Const generics are used to allow easy deconstruction in the callback arguments
///
/// ```rust
///# use trussed::client::{Ed255, CryptoClient};
///# use trussed::types::{Location, Mechanism};
///# use trussed::syscall;
///# use trussed::virt::with_ram_clients;
/// with_ram_clients(["client1", "client2"], |[mut client1, mut client2]| {
/// let key = syscall!(client1.generate_ed255_private_key(Location::Internal)).key;
/// // The clients are distinct
/// assert!(!syscall!(client2.exists(Mechanism::Ed255, key)).exists);
/// })
/// ```
pub fn with_ram_clients<R, F, const N: usize>(client_ids: [&str; N], f: F) -> R
where
F: FnOnce([MultiClient<Ram>; N]) -> R,
{
with_clients(Ram::default(), client_ids, f)
}

pub struct Platform<S: StoreProvider> {
rng: ChaCha8Rng,
_store: S,
ui: UserInterface,
}

impl<S: Syscall> Syscall for Arc<Mutex<S>> {
fn syscall(&mut self) {
self.lock().unwrap().syscall()
}
}

impl<S: StoreProvider> Platform<S> {
pub fn run_client<R>(
self,
client_id: &str,
test: impl FnOnce(ClientImplementation<Service<Self>>) -> R,
) -> R {
pub fn run_client<R>(self, client_id: &str, test: impl FnOnce(Client<S>) -> R) -> R {
let service = Service::new(self);
let client = service.try_into_new_client(client_id, None).unwrap();
test(client)
Expand All @@ -93,7 +140,7 @@ impl<S: StoreProvider> Platform<S> {
client_id: &str,
dispatch: D,
backends: &'static [BackendId<D::BackendId>],
test: impl FnOnce(ClientImplementation<Service<Self, D>, D>) -> R,
test: impl FnOnce(Client<S, D>) -> R,
) -> R {
let mut service = Service::with_dispatch(self, dispatch);
let client = ClientBuilder::new(client_id)
Expand All @@ -103,6 +150,36 @@ impl<S: StoreProvider> Platform<S> {
.build(service);
test(client)
}

pub fn run_clients<R, const N: usize>(
self,
client_ids: [&str; N],
test: impl FnOnce([MultiClient<S>; N]) -> R,
) -> R {
let mut service = Service::new(self);
let prepared_clients =
client_ids.map(|id| ClientBuilder::new(id).prepare(&mut service).unwrap());
let service = Arc::new(Mutex::new(service));
test(prepared_clients.map(|builder| builder.build(service.clone())))
}

/// Using const generics rather than a `Vec` to allow destructuring in the method
pub fn run_clients_with_backends<R, D: Dispatch, const N: usize>(
self,
client_ids: [(&str, &'static [BackendId<D::BackendId>]); N],
dispatch: D,
test: impl FnOnce([MultiClient<S, D>; N]) -> R,
) -> R {
let mut service = Service::with_dispatch(self, dispatch);
let prepared_clients = client_ids.map(|(id, backends)| {
ClientBuilder::new(id)
.backends(backends)
.prepare(&mut service)
.unwrap()
});
let service = Arc::new(Mutex::new(service));
test(prepared_clients.map(|builder| builder.build(service.clone())))
}
}

unsafe impl<S: StoreProvider> platform::Platform for Platform<S> {
Expand Down

0 comments on commit 9763723

Please sign in to comment.