Skip to content

Commit cc2ca52

Browse files
jprendesjprendes
committed
Refactor HostFunction
Signed-off-by: Jorge Prendes <[email protected]>
1 parent 23dcf47 commit cc2ca52

File tree

9 files changed

+307
-282
lines changed

9 files changed

+307
-282
lines changed

‎src/hyperlight_host/src/func/host_functions.rs

Lines changed: 129 additions & 208 deletions
Original file line numberDiff line numberDiff line change
@@ -14,232 +14,153 @@ See the License for the specific language governing permissions and
1414
limitations under the License.
1515
*/
1616

17-
#![allow(non_snake_case)]
1817
use std::sync::{Arc, Mutex};
1918

20-
use hyperlight_common::flatbuffer_wrappers::function_types::ParameterValue;
21-
use tracing::{instrument, Span};
19+
use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnValue};
2220

23-
use super::{HyperlightFunction, SupportedParameterType, SupportedReturnType};
21+
use super::utils::for_each_tuple;
22+
use super::{ParameterTuple, ResultType, SupportedReturnType};
2423
use crate::sandbox::{ExtraAllowedSyscall, UninitializedSandbox};
25-
use crate::HyperlightError::UnexpectedNoOfArguments;
2624
use crate::{log_then_return, new_error, Result};
2725

28-
/// Trait for registering a host function
29-
pub trait HostFunction<R, Args> {
30-
/// Register the host function with the given name in the sandbox.
31-
fn register(&self, sandbox: &mut UninitializedSandbox, name: &str) -> Result<()>;
32-
33-
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
34-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
35-
fn register_with_extra_allowed_syscalls(
36-
&self,
37-
sandbox: &mut UninitializedSandbox,
38-
name: &str,
39-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
40-
) -> Result<()>;
26+
/// A representation of a host function.
27+
/// This is a thin wrapper around a `Fn(Args) -> Result<Output>`.
28+
#[derive(Clone)]
29+
pub struct HostFunction<Output, Args>
30+
where
31+
Args: ParameterTuple,
32+
Output: SupportedReturnType,
33+
{
34+
// This is a thin wrapper around a `Fn(Args) -> Result<Output>`.
35+
// But unlike `Fn` which is a trait, this is a concrete type.
36+
// This allows us to:
37+
// 1. Impose constraints on the function arguments and return type.
38+
// 2. Impose a single function signature.
39+
//
40+
// This second point is important because the `Fn` trait is generic
41+
// over the function arguments (with an associated return type).
42+
// This means that a given type could implement `Fn` for multiple
43+
// function signatures.
44+
// This means we can't do something like:
45+
// ```rust,ignore
46+
// impl<Args, Output, F> SomeTrait for F
47+
// where
48+
// F: Fn(Args) -> Result<Output>,
49+
// { ... }
50+
// ```
51+
// because the concrete type F might implement `Fn` for multiple times,
52+
// and that would means implementing `SomeTrait` multiple times for the
53+
// same type.
54+
55+
// Use Arc in here instead of Box because it's useful in tests and
56+
// presumably in other places to be able to clone a HostFunction and
57+
// use it across different sandboxes.
58+
func: Arc<dyn Fn(Args) -> Result<Output> + Send + Sync + 'static>,
4159
}
4260

43-
/// Tait for types that can be converted into types implementing `HostFunction`.
44-
pub trait IntoHostFunction<R, Args> {
45-
/// Concrete type of the returned host function
46-
type Output: HostFunction<R, Args>;
47-
48-
/// Convert the type into a host function
49-
fn into_host_function(self) -> Self::Output;
61+
pub(crate) struct TypeErasedHostFunction {
62+
func: Box<dyn Fn(Vec<ParameterValue>) -> Result<ReturnValue> + Send + Sync + 'static>,
5063
}
5164

52-
macro_rules! impl_host_function {
53-
(@count) => { 0 };
54-
(@count $P:ident $(, $R:ident)*) => {
55-
impl_host_function!(@count $($R),*) + 1
56-
};
57-
(@impl $($P:ident),*) => {
58-
const _: () = {
59-
impl<R $(, $P)*, F> HostFunction<R, ($($P,)*)> for Arc<Mutex<F>>
60-
where
61-
F: FnMut($($P),*) -> Result<R> + Send + 'static,
62-
$($P: SupportedParameterType + Clone,)*
63-
R: SupportedReturnType,
64-
{
65-
/// Register the host function with the given name in the sandbox.
66-
#[instrument(
67-
err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
68-
)]
69-
fn register(
70-
&self,
71-
sandbox: &mut UninitializedSandbox,
72-
name: &str,
73-
) -> Result<()> {
74-
register_host_function(self.clone(), sandbox, name, None)
75-
}
76-
77-
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
78-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
79-
#[instrument(
80-
err(Debug), skip(self, sandbox, extra_allowed_syscalls),
81-
parent = Span::current(), level = "Trace"
82-
)]
83-
fn register_with_extra_allowed_syscalls(
84-
&self,
85-
sandbox: &mut UninitializedSandbox,
86-
name: &str,
87-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
88-
) -> Result<()> {
89-
register_host_function(self.clone(), sandbox, name, Some(extra_allowed_syscalls))
90-
}
91-
}
92-
93-
impl<R $(, $P)*> HostFunction<R, ($($P,)*)> for &dyn HostFunction<R, ($($P,)*)>
94-
where
95-
$($P: SupportedParameterType + Clone,)*
96-
R: SupportedReturnType,
97-
{
98-
/// Register the host function with the given name in the sandbox.
99-
#[instrument(
100-
err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
101-
)]
102-
fn register(
103-
&self,
104-
sandbox: &mut UninitializedSandbox,
105-
name: &str,
106-
) -> Result<()> {
107-
(**self).register(sandbox, name)
108-
}
109-
110-
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
111-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
112-
#[instrument(
113-
err(Debug), skip(self, sandbox, extra_allowed_syscalls),
114-
parent = Span::current(), level = "Trace"
115-
)]
116-
fn register_with_extra_allowed_syscalls(
117-
&self,
118-
sandbox: &mut UninitializedSandbox,
119-
name: &str,
120-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
121-
) -> Result<()> {
122-
(**self).register_with_extra_allowed_syscalls(sandbox, name, extra_allowed_syscalls)
123-
}
124-
}
125-
126-
impl<R $(, $P)*, F> IntoHostFunction<R, ($($P,)*)> for F
127-
where
128-
F: FnMut($($P),*) -> Result<R> + Send + 'static,
129-
Arc<Mutex<F>>: HostFunction<R, ($($P,)*)>,
130-
{
131-
type Output = Arc<Mutex<F>>;
132-
133-
fn into_host_function(self) -> Self::Output {
134-
Arc::new(Mutex::new(self))
135-
}
136-
}
137-
138-
impl<R $(, $P)*, F> IntoHostFunction<R, ($($P,)*)> for Arc<Mutex<F>>
139-
where
140-
F: FnMut($($P),*) -> Result<R> + Send + 'static,
141-
Arc<Mutex<F>>: HostFunction<R, ($($P,)*)>,
142-
{
143-
type Output = Arc<Mutex<F>>;
65+
impl<Args, Output> HostFunction<Output, Args>
66+
where
67+
Args: ParameterTuple,
68+
Output: SupportedReturnType,
69+
{
70+
/// Call the host function with the given arguments.
71+
pub fn call(&self, args: Args) -> Result<Output> {
72+
(self.func)(args)
73+
}
74+
}
14475

145-
fn into_host_function(self) -> Self::Output {
146-
self
147-
}
148-
}
76+
impl TypeErasedHostFunction {
77+
pub(crate) fn call(&self, args: Vec<ParameterValue>) -> Result<ReturnValue> {
78+
(self.func)(args)
79+
}
80+
}
14981

150-
impl<R $(, $P)*, F> IntoHostFunction<R, ($($P,)*)> for &Arc<Mutex<F>>
151-
where
152-
F: FnMut($($P),*) -> Result<R> + Send + 'static,
153-
Arc<Mutex<F>>: HostFunction<R, ($($P,)*)>,
154-
{
155-
type Output = Arc<Mutex<F>>;
82+
impl<Args, Output> From<HostFunction<Output, Args>> for TypeErasedHostFunction
83+
where
84+
Args: ParameterTuple,
85+
Output: SupportedReturnType,
86+
{
87+
fn from(func: HostFunction<Output, Args>) -> TypeErasedHostFunction {
88+
TypeErasedHostFunction {
89+
func: Box::new(move |args: Vec<ParameterValue>| {
90+
let args = Args::from_value(args)?;
91+
Ok(func.call(args)?.into_value())
92+
}),
93+
}
94+
}
95+
}
15696

157-
fn into_host_function(self) -> Self::Output {
158-
self.clone()
97+
macro_rules! impl_host_function {
98+
([$N:expr] ($($p:ident: $P:ident),*)) => {
99+
/*
100+
// Normally for a `Fn + Send + Sync` we don't need to use a Mutex
101+
// like we do in the case of a `FnMut`.
102+
// However, we can't implement `IntoHostFunction` for `Fn` and `FnMut`
103+
// because `FnMut` is a supertrait of `Fn`.
104+
*/
105+
106+
impl<F, R, $($P),*> From<F> for HostFunction<R::ReturnType, ($($P,)*)>
107+
where
108+
F: FnMut($($P),*) -> R + Send + 'static,
109+
($($P,)*): ParameterTuple,
110+
R: ResultType,
111+
{
112+
fn from(mut func: F) -> HostFunction<R::ReturnType, ($($P,)*)> {
113+
let func = move |($($p,)*): ($($P,)*)| -> Result<R::ReturnType> {
114+
func($($p),*).into_result()
115+
};
116+
let func = Mutex::new(func);
117+
HostFunction {
118+
func: Arc::new(move |args: ($($P,)*)| {
119+
func.try_lock()
120+
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
121+
(args)
122+
})
159123
}
160124
}
125+
}
126+
};
127+
}
161128

162-
impl<R $(, $P)*> IntoHostFunction<R, ($($P,)*)> for &dyn HostFunction<R, ($($P,)*)>
163-
where
164-
R: SupportedReturnType,
165-
$($P: SupportedParameterType + Clone,)*
166-
{
167-
type Output = Self;
129+
for_each_tuple!(impl_host_function);
168130

169-
fn into_host_function(self) -> Self::Output {
170-
self
171-
}
172-
}
131+
pub(crate) fn register_host_function<Args: ParameterTuple, Output: SupportedReturnType>(
132+
func: impl Into<HostFunction<Output, Args>>,
133+
sandbox: &mut UninitializedSandbox,
134+
name: &str,
135+
extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
136+
) -> Result<()> {
137+
let func = func.into().into();
173138

174-
fn register_host_function<T, $($P,)* R>(
175-
self_: Arc<Mutex<T>>,
176-
sandbox: &mut UninitializedSandbox,
177-
name: &str,
178-
extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
179-
) -> Result<()>
180-
where
181-
T: FnMut($($P),*) -> Result<R> + Send + 'static,
182-
$($P: SupportedParameterType + Clone,)*
183-
R: SupportedReturnType,
139+
if let Some(_eas) = extra_allowed_syscalls {
140+
if cfg!(all(feature = "seccomp", target_os = "linux")) {
141+
// Register with extra allowed syscalls
142+
#[cfg(all(feature = "seccomp", target_os = "linux"))]
184143
{
185-
const N: usize = impl_host_function!(@count $($P),*);
186-
let cloned = self_.clone();
187-
let func = Box::new(move |args: Vec<ParameterValue>| {
188-
let ($($P,)*) = match <[ParameterValue; N]>::try_from(args) {
189-
Ok([$($P,)*]) => ($($P::from_value($P)?,)*),
190-
Err(args) => { log_then_return!(UnexpectedNoOfArguments(args.len(), N)); }
191-
};
192-
193-
let result = cloned
194-
.try_lock()
195-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?(
196-
$($P),*
197-
)?;
198-
Ok(result.into_value())
199-
});
200-
201-
if let Some(_eas) = extra_allowed_syscalls {
202-
if cfg!(all(feature = "seccomp", target_os = "linux")) {
203-
// Register with extra allowed syscalls
204-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
205-
{
206-
sandbox
207-
.host_funcs
208-
.try_lock()
209-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
210-
.register_host_function_with_syscalls(
211-
name.to_string(),
212-
HyperlightFunction::new(func),
213-
_eas,
214-
)?;
215-
}
216-
} else {
217-
// Log and return an error
218-
log_then_return!("Extra allowed syscalls are only supported on Linux with seccomp enabled");
219-
}
220-
} else {
221-
// Register without extra allowed syscalls
222-
sandbox
223-
.host_funcs
224-
.try_lock()
225-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
226-
.register_host_function(
227-
name.to_string(),
228-
HyperlightFunction::new(func),
229-
)?;
230-
}
231-
232-
Ok(())
144+
sandbox
145+
.host_funcs
146+
.try_lock()
147+
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
148+
.register_host_function_with_syscalls(name.to_string(), func, _eas)?;
233149
}
234-
};
235-
};
236-
() => {
237-
impl_host_function!(@impl);
238-
};
239-
($P:ident $(, $R:ident)*) => {
240-
impl_host_function!($($R),*);
241-
impl_host_function!(@impl $P $(, $R)*);
242-
};
150+
} else {
151+
// Log and return an error
152+
log_then_return!(
153+
"Extra allowed syscalls are only supported on Linux with seccomp enabled"
154+
);
155+
}
156+
} else {
157+
// Register without extra allowed syscalls
158+
sandbox
159+
.host_funcs
160+
.try_lock()
161+
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
162+
.register_host_function(name.to_string(), func)?;
163+
}
164+
165+
Ok(())
243166
}
244-
245-
impl_host_function!(P1, P2, P3, P4, P5, P6, P7, P8, P9, P10);

0 commit comments

Comments
 (0)