hyperlight-host 0.16.0

/*
Copyright 2025  The Hyperlight Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

use std::collections::HashMap;
use std::io::{IsTerminal, Write};

use hyperlight_common::flatbuffer_wrappers::function_types::{
    ParameterType, ParameterValue, ReturnType, ReturnValue,
};
use hyperlight_common::flatbuffer_wrappers::host_function_definition::HostFunctionDefinition;
use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
use tracing::{Span, instrument};

use crate::HyperlightError::HostFunctionNotFound;
use crate::Result;
use crate::func::host_functions::TypeErasedHostFunction;

#[derive(Default)]
/// A Wrapper around details of functions exposed by the Host
pub struct FunctionRegistry {
    functions_map: HashMap<String, FunctionEntry>,
}

/// A collection of host functions that can be supplied to a sandbox
/// constructor (e.g. [`crate::MultiUseSandbox::from_snapshot`]) to
/// expose host-side functionality to the guest.
///
/// Use [`HostFunctions::default`] to start with the standard
/// `HostPrint` function pre-registered (matches the registry that the
/// regular `UninitializedSandbox` → `evolve()` path constructs), or
/// [`HostFunctions::empty`] to start with an empty registry.
///
/// Add additional host functions via the
/// [`crate::func::Registerable`] trait, just as you would on an
/// `UninitializedSandbox`.
///
/// ```no_run
/// # use hyperlight_host::{HostFunctions, Result};
/// # use hyperlight_host::func::Registerable;
/// # fn example() -> Result<()> {
/// // Default: HostPrint already registered.
/// let mut funcs = HostFunctions::default();
/// funcs.register_host_function("Add", |a: i32, b: i32| Ok(a + b))?;
/// # Ok(())
/// # }
/// ```
pub struct HostFunctions(FunctionRegistry);

impl HostFunctions {
    /// Create an empty `HostFunctions` with no host functions
    /// registered.
    ///
    /// Most callers want [`HostFunctions::default`] instead, which
    /// pre-registers the standard `HostPrint` function. An empty
    /// registry will fail snapshot validation against any snapshot
    /// that captured `HostPrint`, and any guest code that tries to
    /// `printf` into an empty registry will get an EIO from
    /// `write(2)`.
    pub fn empty() -> Self {
        Self(FunctionRegistry::default())
    }

    /// Consume this `HostFunctions` and return the inner registry.
    pub(crate) fn into_inner(self) -> FunctionRegistry {
        self.0
    }

    /// Borrow the inner registry mutably.
    pub(crate) fn inner_mut(&mut self) -> &mut FunctionRegistry {
        &mut self.0
    }

    /// Borrow the inner registry immutably.
    pub(crate) fn inner(&self) -> &FunctionRegistry {
        &self.0
    }
}

impl Default for HostFunctions {
    /// Create a `HostFunctions` pre-populated with the standard
    /// `HostPrint` function (writes UTF-8 strings to the host's
    /// stdout in green).
    ///
    /// This matches the default registry installed by
    /// `UninitializedSandbox::new()`, so a snapshot taken from a
    /// regular sandbox can be loaded with
    /// `MultiUseSandbox::from_snapshot(snap, HostFunctions::default(), None)`
    /// without registering anything else.
    ///
    /// Use [`HostFunctions::empty`] for an empty registry.
    fn default() -> Self {
        Self(FunctionRegistry::with_default_host_print())
    }
}

impl From<&FunctionRegistry> for HostFunctionDetails {
    fn from(registry: &FunctionRegistry) -> Self {
        let host_functions = registry
            .functions_map
            .iter()
            .map(|(name, entry)| HostFunctionDefinition {
                function_name: name.clone(),
                parameter_types: Some(entry.parameter_types.to_vec()),
                return_type: entry.return_type,
            })
            .collect();

        HostFunctionDetails {
            host_functions: Some(host_functions),
        }
    }
}

pub struct FunctionEntry {
    pub function: TypeErasedHostFunction,
    pub parameter_types: &'static [ParameterType],
    pub return_type: ReturnType,
}

impl FunctionRegistry {
    /// Register a host function with the sandbox.
    #[instrument(skip_all, parent = Span::current(), level = "Trace")]
    pub(crate) fn register_host_function(&mut self, name: String, func: FunctionEntry) {
        self.functions_map.insert(name, func);
    }

    /// Return the registered signature for `name`.
    pub(crate) fn function_signature(
        &self,
        name: &str,
    ) -> Option<(&'static [ParameterType], ReturnType)> {
        self.functions_map
            .get(name)
            .map(|entry| (entry.parameter_types, entry.return_type))
    }

    /// Create a `FunctionRegistry` pre-populated with the default
    /// `HostPrint` function (writes to stdout with green text).
    pub(crate) fn with_default_host_print() -> Self {
        use crate::func::host_functions::HostFunction;
        use crate::func::{ParameterTuple, SupportedReturnType};

        let mut registry = Self::default();
        let hf: HostFunction<i32, (String,)> = default_writer_func.into();
        let entry = FunctionEntry {
            function: hf.into(),
            parameter_types: <(String,)>::TYPE,
            return_type: <i32 as SupportedReturnType>::TYPE,
        };
        registry.register_host_function("HostPrint".to_string(), entry);
        registry
    }

    /// Assuming a host function called `"HostPrint"` exists, and takes a
    /// single string parameter, call it with the given `msg` parameter.
    ///
    /// Return `Ok` if the function was found and was of the right signature,
    /// and `Err` otherwise.
    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
    #[allow(dead_code)]
    pub(super) fn host_print(&mut self, msg: String) -> Result<i32> {
        let res = self.call_host_func_impl("HostPrint", vec![ParameterValue::String(msg)])?;
        res.try_into()
            .map_err(|_| HostFunctionNotFound("HostPrint".to_string()))
    }
    /// From the set of registered host functions, attempt to get the one
    /// named `name`. If it exists, call it with the given arguments list
    /// `args` and return its result.
    ///
    /// Return `Err` if no such function exists,
    /// its parameter list doesn't match `args`, or there was another error
    /// getting, configuring or calling the function.
    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
    pub(super) fn call_host_function(
        &self,
        name: &str,
        args: Vec<ParameterValue>,
    ) -> Result<ReturnValue> {
        self.call_host_func_impl(name, args)
    }

    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
    fn call_host_func_impl(&self, name: &str, args: Vec<ParameterValue>) -> Result<ReturnValue> {
        let FunctionEntry {
            function,
            parameter_types: _,
            return_type: _,
        } = self
            .functions_map
            .get(name)
            .ok_or_else(|| HostFunctionNotFound(name.to_string()))?;

        // Make the host function call
        crate::metrics::maybe_time_and_emit_host_call(name, || function.call(args))
    }
}

/// The default writer function is to write to stdout with green text.
#[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
fn default_writer_func(s: String) -> Result<i32> {
    match std::io::stdout().is_terminal() {
        false => {
            print!("{}", s);
            Ok(s.len() as i32)
        }
        true => {
            let mut stdout = StandardStream::stdout(ColorChoice::Auto);
            let mut color_spec = ColorSpec::new();
            color_spec.set_fg(Some(Color::Green));
            stdout.set_color(&color_spec)?;
            stdout.write_all(s.as_bytes())?;
            stdout.reset()?;
            Ok(s.len() as i32)
        }
    }
}