hyperlight-host 0.17.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.
use std::io::Write;
use std::sync::{Arc, Mutex};

use fallible_iterator::FallibleIterator;
use framehop::Unwinder;

use crate::hypervisor::regs::CommonRegisters;
#[cfg(not(unshared_snapshot_mem))]
use crate::mem::layout::ReadableSharedMemory;
use crate::mem::layout::SandboxMemoryLayout;
use crate::mem::mgr::SandboxMemoryManager;
use crate::mem::shared_mem::HostSharedMemory;
use crate::sandbox::outb::HandleOutbError;
use crate::{Result, new_error};

/// The type of trace frame being recorded.
/// This is used to identify the type of frame being recorded in the trace file.
enum TraceFrameType {
    /// A frame that records a memory allocation.
    MemAlloc = 2,
    /// A frame that records a memory free.
    MemFree = 3,
}

/// This structure handles the memory profiling trace information.
pub(crate) struct MemTraceInfo {
    /// The epoch against which trace events are timed; at least as
    /// early as the creation of the sandbox being traced.
    epoch: std::time::Instant,
    /// The file to which the trace is being written
    pub file: Arc<Mutex<std::fs::File>>,
    /// The unwind information for the current guest
    #[allow(dead_code)]
    pub unwind_module: Arc<dyn crate::mem::exe::UnwindInfo>,
    /// The framehop unwinder for the current guest
    pub unwinder: framehop::x86_64::UnwinderX86_64<Vec<u8>>,
    /// The framehop cache
    pub unwind_cache: Arc<Mutex<framehop::x86_64::CacheX86_64>>,
}

impl MemTraceInfo {
    pub fn new(unwind_module: Arc<dyn crate::mem::exe::UnwindInfo>) -> Result<Self> {
        let mut path = std::env::current_dir()?;
        path.push("trace");

        // create directory if it does not exist
        if !path.exists() {
            std::fs::create_dir(&path)?;
        }
        path.push(uuid::Uuid::new_v4().to_string());
        path.set_extension("trace");

        tracing::info!("Creating trace file at: {}", path.display());
        println!("Creating trace file at: {}", path.display());

        let hash = unwind_module.hash();
        let (unwinder, unwind_cache) = {
            let mut unwinder = framehop::x86_64::UnwinderX86_64::new();
            unwinder.add_module(unwind_module.clone().as_module());
            let cache = framehop::x86_64::CacheX86_64::new();
            (unwinder, Arc::new(Mutex::new(cache)))
        };

        let ret = Self {
            epoch: std::time::Instant::now(),
            file: Arc::new(Mutex::new(std::fs::File::create_new(path)?)),
            unwind_module,
            unwinder,
            unwind_cache,
        };

        /* write a frame identifying the binary */
        ret.record_trace_frame(ret.epoch, 0, |f| {
            let _ = f.write_all(hash.as_bytes());
        })?;

        Ok(ret)
    }

    fn unwind(
        &self,
        regs: &CommonRegisters,
        mem_mgr: &SandboxMemoryManager<HostSharedMemory>,
    ) -> Result<Vec<u64>> {
        let mut read_stack = |addr| {
            let mut buf: [u8; 8] = [0u8; 8];
            mem_mgr
                .shared_mem
                .copy_to_slice(
                    &mut buf,
                    (addr - SandboxMemoryLayout::BASE_ADDRESS as u64) as usize,
                )
                .map_err(|_| ())?;
            Ok(u64::from_ne_bytes(buf))
        };
        let mut cache = self
            .unwind_cache
            .try_lock()
            .map_err(|e| new_error!("could not lock unwinder cache {}\n", e))?;
        let iter = self.unwinder.iter_frames(
            regs.rip,
            framehop::x86_64::UnwindRegsX86_64::new(regs.rip, regs.rsp, regs.rbp),
            &mut *cache,
            &mut read_stack,
        );
        iter.map(|f| Ok(f.address() - mem_mgr.layout.get_guest_code_address() as u64))
            .collect()
            .map_err(|e| new_error!("couldn't unwind: {}", e))
    }

    fn write_stack(&self, out: &mut std::fs::File, stack: &[u64]) {
        let _ = out.write_all(&stack.len().to_ne_bytes());
        for frame in stack {
            let _ = out.write_all(&frame.to_ne_bytes());
        }
    }

    fn record_trace_frame<F: FnOnce(&mut std::fs::File)>(
        &self,
        start_instant: std::time::Instant,
        frame_id: u64,
        write_frame: F,
    ) -> Result<()> {
        let Ok(mut out) = self.file.lock() else {
            return Ok(());
        };
        // frame structure:
        // 16 bytes timestamp
        let now = std::time::Instant::now().saturating_duration_since(start_instant);
        let _ = out.write_all(&now.as_micros().to_ne_bytes());
        // 8 bytes frame type id
        let _ = out.write_all(&frame_id.to_ne_bytes());
        // frame data
        write_frame(&mut out);
        Ok(())
    }

    fn handle_trace(
        &self,
        regs: &CommonRegisters,
        mem_mgr: &SandboxMemoryManager<HostSharedMemory>,
        trace_identifier: TraceFrameType,
    ) -> std::result::Result<(), HandleOutbError> {
        let Ok(stack) = self.unwind(regs, mem_mgr) else {
            return Ok(());
        };

        let amt = regs.rax;
        let ptr = regs.rcx;

        match trace_identifier {
            TraceFrameType::MemAlloc => self
                .record_trace_frame(self.epoch, trace_identifier as u64, |f| {
                    let _ = f.write_all(&ptr.to_ne_bytes());
                    let _ = f.write_all(&amt.to_ne_bytes());
                    self.write_stack(f, &stack);
                })
                .map_err(|e| HandleOutbError::MemProfile(e.to_string())),
            // The MemFree case does not expect an amount, only a pointer
            TraceFrameType::MemFree => self
                .record_trace_frame(self.epoch, trace_identifier as u64, |f| {
                    let _ = f.write_all(&ptr.to_ne_bytes());
                    self.write_stack(f, &stack);
                })
                .map_err(|e| HandleOutbError::MemProfile(e.to_string())),
        }
    }

    #[inline(always)]
    pub(crate) fn handle_trace_mem_alloc(
        &self,
        regs: &CommonRegisters,
        mem_mgr: &SandboxMemoryManager<HostSharedMemory>,
    ) -> std::result::Result<(), HandleOutbError> {
        self.handle_trace(regs, mem_mgr, TraceFrameType::MemAlloc)
    }

    #[inline(always)]
    pub(crate) fn handle_trace_mem_free(
        &self,
        regs: &CommonRegisters,
        mem_mgr: &SandboxMemoryManager<HostSharedMemory>,
    ) -> std::result::Result<(), HandleOutbError> {
        self.handle_trace(regs, mem_mgr, TraceFrameType::MemFree)
    }
}