Skip to content

Commit 0454d5f

Browse files
jprendesjprendes
committed
Unify all the HostFunctionXX traits into a single HostFunction trait
Signed-off-by: Jorge Prendes <[email protected]>
1 parent 3d455ca commit 0454d5f

File tree

1 file changed

+46
-157
lines changed

1 file changed

+46
-157
lines changed

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

Lines changed: 46 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ use std::sync::{Arc, Mutex};
2020
use hyperlight_common::flatbuffer_wrappers::function_types::{
2121
ParameterType, ParameterValue, ReturnType,
2222
};
23-
use paste::paste;
2423
use tracing::{instrument, Span};
2524

2625
use super::{HyperlightFunction, SupportedParameterType, SupportedReturnType};
@@ -54,145 +53,39 @@ impl HostFunctionDefinition {
5453
}
5554
}
5655

57-
macro_rules! host_function {
58-
// Special case for zero parameters
59-
(0) => {
60-
paste! {
61-
/// Trait for registering a host function with zero parameters.
62-
pub trait HostFunction0<'a, R: SupportedReturnType<R>> {
63-
/// Register the host function with the given name in the sandbox.
64-
fn register(
65-
&self,
66-
sandbox: &mut UninitializedSandbox,
67-
name: &str,
68-
) -> Result<()>;
69-
70-
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
71-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
72-
fn register_with_extra_allowed_syscalls(
73-
&self,
74-
sandbox: &mut UninitializedSandbox,
75-
name: &str,
76-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
77-
) -> Result<()>;
78-
}
79-
80-
impl<'a, T, R> HostFunction0<'a, R> for Arc<Mutex<T>>
81-
where
82-
T: FnMut() -> Result<R> + Send + 'static,
83-
R: SupportedReturnType<R>,
84-
{
85-
#[instrument(
86-
err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
87-
)]
88-
fn register(
89-
&self,
90-
sandbox: &mut UninitializedSandbox,
91-
name: &str,
92-
) -> Result<()> {
93-
register_host_function_0(self.clone(), sandbox, name, None)
94-
}
95-
96-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
97-
#[instrument(
98-
err(Debug), skip(self, sandbox, extra_allowed_syscalls),
99-
parent = Span::current(), level = "Trace"
100-
)]
101-
fn register_with_extra_allowed_syscalls(
102-
&self,
103-
sandbox: &mut UninitializedSandbox,
104-
name: &str,
105-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
106-
) -> Result<()> {
107-
register_host_function_0(self.clone(), sandbox, name, Some(extra_allowed_syscalls))
108-
}
109-
}
110-
111-
fn register_host_function_0<T, R>(
112-
self_: Arc<Mutex<T>>,
113-
sandbox: &mut UninitializedSandbox,
114-
name: &str,
115-
extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
116-
) -> Result<()>
117-
where
118-
T: FnMut() -> Result<R> + Send + 'static,
119-
R: SupportedReturnType<R>,
120-
{
121-
let cloned = self_.clone();
122-
let func = Box::new(move |_: Vec<ParameterValue>| {
123-
let result = cloned
124-
.try_lock()
125-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?()?;
126-
Ok(result.get_hyperlight_value())
127-
});
128-
129-
if let Some(_eas) = extra_allowed_syscalls {
130-
if cfg!(all(feature = "seccomp", target_os = "linux")) {
131-
// Register with extra allowed syscalls
132-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
133-
{
134-
sandbox
135-
.host_funcs
136-
.try_lock()
137-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
138-
.register_host_function_with_syscalls(
139-
&HostFunctionDefinition::new(name.to_string(), None, R::get_hyperlight_type()),
140-
HyperlightFunction::new(func),
141-
_eas,
142-
)?;
143-
}
144-
} else {
145-
// Log and return an error
146-
log_then_return!("Extra allowed syscalls are only supported on Linux with seccomp enabled");
147-
}
148-
} else {
149-
// Register without extra allowed syscalls
150-
sandbox
151-
.host_funcs
152-
.try_lock()
153-
.map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
154-
.register_host_function(
155-
&HostFunctionDefinition::new(name.to_string(), None, R::get_hyperlight_type()),
156-
HyperlightFunction::new(func),
157-
)?;
158-
}
56+
/// Trait for registering a host function
57+
pub trait HostFunction<R, Args> {
58+
/// Register the host function with the given name in the sandbox.
59+
fn register(
60+
&self,
61+
sandbox: &mut UninitializedSandbox,
62+
name: &str,
63+
) -> Result<()>;
64+
65+
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
66+
#[cfg(all(feature = "seccomp", target_os = "linux"))]
67+
fn register_with_extra_allowed_syscalls(
68+
&self,
69+
sandbox: &mut UninitializedSandbox,
70+
name: &str,
71+
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
72+
) -> Result<()>;
73+
}
15974

160-
Ok(())
161-
}
162-
}
75+
macro_rules! impl_host_function {
76+
(@count) => { 0 };
77+
(@count $P:ident $(, $R:ident)*) => {
78+
impl_host_function!(@count $($R),*) + 1
16379
};
164-
// General case for one or more parameters
165-
($N:expr, $($P:ident),+) => {
166-
paste! {
167-
/// Trait for registering a host function with $N parameters.
168-
pub trait [<HostFunction $N>]<'a, $($P,)* R>
80+
(@impl $($P:ident),*) => {
81+
const _: () = {
82+
impl<'a, R $(, $P)*, F> HostFunction<R, ($($P,)*)> for Arc<Mutex<F>>
16983
where
84+
F: FnMut($($P),*) -> Result<R> + Send + 'static,
17085
$($P: SupportedParameterType<$P> + Clone + 'a,)*
17186
R: SupportedReturnType<R>,
17287
{
17388
/// Register the host function with the given name in the sandbox.
174-
fn register(
175-
&self,
176-
sandbox: &mut UninitializedSandbox,
177-
name: &str,
178-
) -> Result<()>;
179-
180-
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
181-
#[cfg(all(feature = "seccomp", target_os = "linux"))]
182-
fn register_with_extra_allowed_syscalls(
183-
&self,
184-
sandbox: &mut UninitializedSandbox,
185-
name: &str,
186-
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
187-
) -> Result<()>;
188-
}
189-
190-
impl<'a, T, $($P,)* R> [<HostFunction $N>]<'a, $($P,)* R> for Arc<Mutex<T>>
191-
where
192-
T: FnMut($($P),*) -> Result<R> + Send + 'static,
193-
$($P: SupportedParameterType<$P> + Clone + 'a,)*
194-
R: SupportedReturnType<R>,
195-
{
19689
#[instrument(
19790
err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
19891
)]
@@ -201,9 +94,10 @@ macro_rules! host_function {
20194
sandbox: &mut UninitializedSandbox,
20295
name: &str,
20396
) -> Result<()> {
204-
[<register_host_function_ $N>](self.clone(), sandbox, name, None)
97+
register_host_function(self.clone(), sandbox, name, None)
20598
}
20699

100+
/// Register the host function with the given name in the sandbox, allowing extra syscalls.
207101
#[cfg(all(feature = "seccomp", target_os = "linux"))]
208102
#[instrument(
209103
err(Debug), skip(self, sandbox, extra_allowed_syscalls),
@@ -215,11 +109,11 @@ macro_rules! host_function {
215109
name: &str,
216110
extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
217111
) -> Result<()> {
218-
[<register_host_function_ $N>](self.clone(), sandbox, name, Some(extra_allowed_syscalls))
112+
register_host_function(self.clone(), sandbox, name, Some(extra_allowed_syscalls))
219113
}
220114
}
221115

222-
fn [<register_host_function_ $N>]<'a, T, $($P,)* R>(
116+
fn register_host_function<'a, T, $($P,)* R>(
223117
self_: Arc<Mutex<T>>,
224118
sandbox: &mut UninitializedSandbox,
225119
name: &str,
@@ -231,15 +125,13 @@ macro_rules! host_function {
231125
R: SupportedReturnType<R>,
232126
{
233127
let cloned = self_.clone();
234-
let func = Box::new(move |args: Vec<ParameterValue>| {
235-
if args.len() != $N {
236-
log_then_return!(UnexpectedNoOfArguments(args.len(), $N));
237-
}
238-
239-
let mut args_iter = args.into_iter();
240-
$(
241-
let $P = $P::get_inner(args_iter.next().unwrap())?;
242-
)*
128+
let func = Box::new(move |mut args: Vec<ParameterValue>| {
129+
let ($($P,)*) = match &mut args[..] {
130+
[$($P,)*] => ($($P::get_inner(core::mem::replace($P, ParameterValue::Int(0)))?,)*),
131+
_ => {
132+
log_then_return!(UnexpectedNoOfArguments(args.len(), impl_host_function!(@count $($P),*)));
133+
}
134+
};
243135

244136
let result = cloned
245137
.try_lock()
@@ -292,18 +184,15 @@ macro_rules! host_function {
292184

293185
Ok(())
294186
}
295-
}
187+
};
188+
};
189+
() => {
190+
impl_host_function!(@impl);
191+
};
192+
($P:ident $(, $R:ident)*) => {
193+
impl_host_function!($($R),*);
194+
impl_host_function!(@impl $P $(, $R)*);
296195
};
297196
}
298197

299-
host_function!(0);
300-
host_function!(1, P1);
301-
host_function!(2, P1, P2);
302-
host_function!(3, P1, P2, P3);
303-
host_function!(4, P1, P2, P3, P4);
304-
host_function!(5, P1, P2, P3, P4, P5);
305-
host_function!(6, P1, P2, P3, P4, P5, P6);
306-
host_function!(7, P1, P2, P3, P4, P5, P6, P7);
307-
host_function!(8, P1, P2, P3, P4, P5, P6, P7, P8);
308-
host_function!(9, P1, P2, P3, P4, P5, P6, P7, P8, P9);
309-
host_function!(10, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10);
198+
impl_host_function!(P1, P2, P3, P4, P5, P6, P7, P8, P9, P10);

0 commit comments

Comments
 (0)