hyperlight-host 0.17.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.

use std::collections::HashMap;
use std::time::{Duration, Instant, SystemTime};

use hyperlight_common::flatbuffer_wrappers::guest_trace_data::{
    EventKeyValue, EventsBatchDecoder, EventsDecoder, GuestEvent, MAX_TRACE_DATA_SIZE,
};
use hyperlight_common::outb::OutBAction;
use opentelemetry::global::BoxedSpan;
use opentelemetry::trace::{Span as _, TraceContextExt, Tracer as _};
use opentelemetry::{Context, KeyValue, global};
use tracing::span::{EnteredSpan, Span};
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::hypervisor::regs::CommonRegisters;
use crate::mem::mgr::SandboxMemoryManager;
use crate::mem::shared_mem::HostSharedMemory;
use crate::{Result, new_error};

/// Type that helps get the data from the guest provided the registers and memory access
struct EventsBatch {
    events: Vec<GuestEvent>,
}

impl EventsBatch {
    /// Extract a batch of guest trace events from guest memory.
    ///
    /// The guest passes the trace data pointer as a Guest Virtual Address (GVA)
    /// in register r9. With Copy-on-Write enabled, this GVA may not be
    /// identity-mapped to its physical address, so we walk the guest page
    /// tables to translate GVA → GPA before reading the data.
    ///
    /// # Arguments
    /// * `regs` - The guest registers (r8 = magic, r9 = GVA pointer, r10 = length)
    /// * `mem_mgr` - The sandbox memory manager with access to shared and scratch memory
    /// * `root_pt` - The root page table physical address (CR3) for GVA translation
    fn from_regs(
        regs: &CommonRegisters,
        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
        root_pt: u64,
    ) -> Result<Self> {
        let magic_no = regs.r8;
        let trace_data_gva = regs.r9;
        let trace_data_len = regs.r10 as usize;

        // Validate the magic number to ensure the guest is providing trace data
        if magic_no != OutBAction::TraceBatch as u64 {
            return Err(new_error!("A TraceBatch is not present"));
        }

        // Validate the length to prevent reading excessive memory
        if trace_data_len == 0 || trace_data_len > MAX_TRACE_DATA_SIZE {
            return Err(new_error!("Invalid TraceBatch length: {}", trace_data_len));
        }

        // Read the trace data from guest memory by walking the page tables
        // to translate the GVA to physical addresses. This is necessary
        // because with CoW, guest virtual pages are backed by physical
        // pages in the scratch region rather than being identity-mapped.
        let buf = mem_mgr.read_guest_memory_by_gva(trace_data_gva, trace_data_len, root_pt)?;

        let events = EventsBatchDecoder {}.decode(&buf).map_err(|e| {
            tracing::error!("Failed to deserialize guest trace events: {:?}", e);
            new_error!("Failed to deserialize guest trace events: {:?}", e)
        })?;

        Ok(EventsBatch { events })
    }
}

/// This structure handles the guest tracing information.
pub struct TraceContext {
    host_spans: Vec<EnteredSpan>,
    guest_spans: HashMap<u64, BoxedSpan>,
    in_host_call: bool,

    // Lazily initialized members
    start_wall: Option<SystemTime>,
    /// The epoch at which the call into the guest started, if it has started.
    /// This is used to calculate the time spent in the guest relative to the
    /// time when the call into the guest was first made.
    start_instant: Option<Instant>,
    /// The start guest time, in TSC cycles, for the current guest measured on the host.
    /// It contains the TSC value recorded on the host before a call is made into the guest.
    /// This is used to calculate the TSC frequency which is the same on the host and guest.
    /// The TSC frequency is used to convert TSC values to timestamps in the trace.
    /// **NOTE**: This is only used until the TSC frequency is calculated, when the first
    /// records are received.
    start_tsc: Option<u64>,
    /// The frequency of the timestamp counter.
    tsc_freq: Option<u64>,
    current_parent_ctx: Option<Context>,
}

impl TraceContext {
    /// Initialize with current context
    pub fn new() -> Self {
        if !hyperlight_guest_tracing::invariant_tsc::has_invariant_tsc() {
            // If the platform does not support invariant TSC, warn the user.
            // On Azure nested virtualization, the TSC invariant bit is not correctly reported, this is a known issue.
            tracing::warn!(
                "Invariant TSC is not supported on this platform, trace timestamps may be inaccurate"
            );
        }

        let current_ctx = Span::current().context();

        let span = tracing::info_span!("call-to-guest");
        let _ = span.set_parent(current_ctx);
        let entered = span.entered();

        Self {
            host_spans: vec![entered],
            guest_spans: HashMap::new(),
            in_host_call: false,

            start_wall: None,
            start_instant: None,
            start_tsc: None,
            tsc_freq: None,
            current_parent_ctx: None,
        }
    }

    /// Calculate the frequency of the TimeStamp Counter.
    /// This is done by:
    /// - first reading a timestamp and an `Instant`
    /// - secondly reading another timestamp and `Instant`
    /// - calculate the frequency based on the `Duration` between
    ///   the two `Instant`s read.
    fn calculate_tsc_freq(&mut self) -> Result<()> {
        let (start, start_time) = match (self.start_tsc.as_ref(), self.start_instant.as_ref()) {
            (Some(start), Some(start_time)) => (*start, *start_time),
            _ => {
                // If the guest start TSC and time are not set, we use the current time and TSC.
                // This is not ideal, but it allows us to calculate the TSC frequency without
                // failing.
                // This is a fallback mechanism to ensure that we can still calculate, however it
                // should be noted that this may lead to inaccuracies in the TSC frequency.
                // The start time should be already set before running the guest for each sandbox.
                tracing::error!(
                    "Guest start TSC and time are not set. Calculating TSC frequency will use current time and TSC."
                );
                (
                    hyperlight_guest_tracing::invariant_tsc::read_tsc(),
                    std::time::Instant::now(),
                )
            }
        };

        let end_time = std::time::Instant::now();
        let end = hyperlight_guest_tracing::invariant_tsc::read_tsc();

        let elapsed = end_time.duration_since(start_time).as_secs_f64();
        let tsc_freq = ((end - start) as f64 / elapsed) as u64;

        tracing::info!("Calculated TSC frequency: {} Hz", tsc_freq);
        self.tsc_freq = Some(tsc_freq);

        Ok(())
    }

    /// Calculate timestamp relative to wall time stored on host
    fn calculate_guest_time_relative_to_host(
        &self,
        guest_start_tsc: u64,
        tsc: u64,
    ) -> Result<SystemTime> {
        // Should never fail as it is extracted after it is set
        let tsc_freq = self.tsc_freq.ok_or(new_error!("TSC frequency not set"))?;

        // Number of cycles relative to guest start
        let rel_cycles = tsc.saturating_sub(guest_start_tsc);

        // Number of micro seconds from guest start to `tsc` argument
        let rel_start_us = rel_cycles as f64 / tsc_freq as f64 * 1_000_000f64;

        // Final timestamp is calculated by:
        // - starting from the wall time when the sandbox was created
        // - adding the Duration to the guest start
        // - adding the Duration from the guest start to the provided `tsc`
        Ok(self.start_wall.ok_or(new_error!("start_wall not set"))?
            + Duration::from_micros(rel_start_us as u64))
    }

    /// Check if the registers indicate that there is trace data to be handled.
    pub fn has_trace_data(&self, regs: &CommonRegisters) -> bool {
        regs.r8 == OutBAction::TraceBatch as u64
    }

    pub fn handle_trace(
        &mut self,
        regs: &CommonRegisters,
        mem_mgr: &mut SandboxMemoryManager<HostSharedMemory>,
        root_pt: u64,
    ) -> Result<()> {
        // Get the guest sent info
        let trace_batch = EventsBatch::from_regs(regs, mem_mgr, root_pt)?;

        self.handle_trace_impl(trace_batch.events)
    }

    fn handle_trace_impl(&mut self, events: Vec<GuestEvent>) -> Result<()> {
        let tracer = global::tracer("guest-tracer");

        // Stack to keep track of open spans
        let mut spans_stack = vec![];

        // Process each event
        for ev in events.into_iter() {
            match ev {
                GuestEvent::GuestStart { tsc } => {
                    // Move to GuestStart
                    if self.tsc_freq.is_none() {
                        self.calculate_tsc_freq()?;
                    }
                    self.start_tsc = Some(tsc);
                }
                GuestEvent::EditSpan { id, fields } => {
                    // Edit existing span attributes
                    if let Some(span) = self.guest_spans.get_mut(&id) {
                        for EventKeyValue { key, value } in fields.iter() {
                            span.set_attribute(KeyValue::new(
                                key.as_str().to_string(),
                                value.as_str().to_string(),
                            ));
                        }
                    } else {
                        tracing::warn!("Tried to edit non-existing guest span with id {}", id);
                    }
                }
                GuestEvent::OpenSpan {
                    id,
                    parent_id,
                    name,
                    target,
                    tsc,
                    fields,
                } => {
                    let start_tsc = self.start_tsc.ok_or(new_error!(
                        "Guest start TSC not set before opening guest span"
                    ))?;
                    // Calculate start timestamp
                    let start_ts = self.calculate_guest_time_relative_to_host(start_tsc, tsc)?;

                    // Determine parent context
                    // Priority:
                    // 1. If parent_id is set and found in guest_spans, use that
                    // 2. If current_parent_ctx is set, use that
                    // 3. Otherwise, use the current span context
                    let parent_ctx = if let Some(parent_id) = parent_id {
                        if let Some(span) = self.guest_spans.get(&parent_id) {
                            Context::new().with_remote_span_context(span.span_context().clone())
                        } else if let Some(parent_ctx) = self.current_parent_ctx.as_ref() {
                            parent_ctx.clone()
                        } else {
                            Span::current().context().clone()
                        }
                    } else if let Some(parent_ctx) = self.current_parent_ctx.as_ref() {
                        parent_ctx.clone()
                    } else {
                        Span::current().context().clone()
                    };

                    // Create the span with calculated start time
                    let mut sb = tracer
                        .span_builder(name.to_string())
                        .with_start_time(start_ts);
                    // Set target attribute
                    sb.attributes = Some(vec![KeyValue::new("target", target.to_string())]);

                    // Attach to parent context
                    let mut span = sb.start_with_context(&tracer, &parent_ctx);

                    // Set attributes from fields
                    for EventKeyValue { key, value } in fields.iter() {
                        span.set_attribute(KeyValue::new(
                            key.as_str().to_string(),
                            value.as_str().to_string(),
                        ));
                    }

                    // Store the span
                    self.guest_spans.insert(id, span);
                    spans_stack.push(id);
                }
                GuestEvent::CloseSpan { id, tsc } => {
                    let start_tsc = self.start_tsc.ok_or(new_error!(
                        "Guest start TSC not set before opening guest span"
                    ))?;
                    // Remove the span and end it
                    if let Some(mut span) = self.guest_spans.remove(&id) {
                        let end_ts = self.calculate_guest_time_relative_to_host(start_tsc, tsc)?;
                        span.end_with_timestamp(end_ts);

                        // The span ids should be closed in order
                        if let Some(stack_id) = spans_stack.pop()
                            && stack_id != id
                        {
                            tracing::warn!("Guest span with id {} closed out of order", id);
                        }
                    } else {
                        tracing::warn!("Tried to close non-existing guest span with id {}", id);
                    }
                }
                GuestEvent::LogEvent {
                    parent_id,
                    name,
                    tsc,
                    fields,
                } => {
                    let start_tsc = self.start_tsc.ok_or(new_error!(
                        "Guest start TSC not set before opening guest span"
                    ))?;
                    let ts = self.calculate_guest_time_relative_to_host(start_tsc, tsc)?;

                    // Add the event to the parent span
                    // It should always have a parent span
                    if let Some(span) = self.guest_spans.get_mut(&parent_id) {
                        let attributes: Vec<KeyValue> = fields
                            .into_iter()
                            .map(|EventKeyValue { key, value }| KeyValue::new(key, value))
                            .collect();
                        span.add_event_with_timestamp(name.to_string(), ts, attributes);
                    } else {
                        tracing::warn!(
                            "Tried to add event to non-existing guest span with id {}",
                            parent_id
                        );
                    }
                }
            }
        }

        // Set the current active span context as the last span in the stack because we want
        // to create a host span that is a child of the last active guest span
        if let Some(span) = spans_stack.pop().and_then(|id| self.guest_spans.get(&id)) {
            // Set as current active span
            let ctx = Context::current().with_remote_span_context(span.span_context().clone());
            self.new_host_trace(ctx);
        };

        Ok(())
    }

    pub(crate) fn setup_guest_trace(&mut self, ctx: Context) {
        if self.start_instant.is_none() {
            crate::debug!("Guest Start Epoch set");
            self.start_wall = Some(SystemTime::now());
            self.start_tsc = Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
            self.start_instant = Some(std::time::Instant::now());
        }
        self.current_parent_ctx = Some(ctx);
    }

    pub fn new_host_trace(&mut self, ctx: Context) {
        let span = tracing::info_span!("call-to-host");
        let _ = span.set_parent(ctx);
        let entered = span.entered();
        self.host_spans.push(entered);
        self.in_host_call = true;
    }

    pub fn end_host_trace(&mut self) {
        if self.in_host_call
            && let Some(entered) = self.host_spans.pop()
        {
            entered.exit();
        }
    }
}

impl Drop for TraceContext {
    fn drop(&mut self) {
        for (k, mut v) in self.guest_spans.drain() {
            v.end();
            tracing::debug!("Dropped guest span with id {}", k);
        }
        while let Some(entered) = self.host_spans.pop() {
            entered.exit();
        }
    }
}

#[cfg(test)]
mod tests {
    use hyperlight_common::flatbuffer_wrappers::guest_trace_data::{EventKeyValue, GuestEvent};

    use super::*;

    fn create_dummy_trace_context() -> TraceContext {
        let mut trace_ctx = TraceContext::new();
        // Set TSC frequency to avoid calculating it
        trace_ctx.tsc_freq = Some(3_200_000_000); // 3.2 GHz
        // Set start wall time and Instant
        trace_ctx.start_wall = Some(SystemTime::now() - Duration::from_secs(1));
        trace_ctx.start_instant = Some(Instant::now() - Duration::from_secs(1));

        trace_ctx
    }

    fn create_open_span(
        id: u64,
        parent_id: Option<u64>,
        name_str: &str,
        target_str: &str,
        start_tsc: u64,
        fields: Vec<EventKeyValue>,
    ) -> GuestEvent {
        GuestEvent::OpenSpan {
            id,
            parent_id,
            name: String::from(name_str),
            target: String::from(target_str),
            tsc: start_tsc,
            fields,
        }
    }

    fn create_close_span(id: u64, end_tsc: u64) -> GuestEvent {
        GuestEvent::CloseSpan { id, tsc: end_tsc }
    }

    fn create_log_event(
        parent_id: u64,
        tsc: u64,
        name_str: &str,
        fields: Vec<EventKeyValue>,
    ) -> GuestEvent {
        GuestEvent::LogEvent {
            parent_id,
            name: String::from(name_str),
            tsc,
            fields,
        }
    }

    #[test]
    fn test_guest_trace_context_creation() {
        let trace_ctx = TraceContext::new();
        assert!(trace_ctx.host_spans.len() == 1);
        assert!(trace_ctx.guest_spans.is_empty());
    }

    /// Test handling a batch with no spans or events.
    #[test]
    fn test_guest_trace_empty_trace_batch() {
        let mut trace_ctx = TraceContext::new();

        let events = vec![];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.is_empty());
        assert!(trace_ctx.host_spans.len() == 1);
    }

    /// Test handling a batch with one span and no events.
    /// The span is not closed.
    #[test]
    fn test_guest_trace_single_span() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "test-span", "test-target", 2000, vec![]),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.len() == 1);
        // The active host span is new because a new guest span was created
        assert!(trace_ctx.host_spans.len() == 2);
    }

    /// Test handling a batch with one span that is closed.
    /// The span is closed.
    #[test]
    fn test_guest_trace_single_closed_span() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "test-span", "test-target", 2000, vec![]),
            create_close_span(1, 2500),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.is_empty());
        // The active host span is the same as before because no new guest span was created
        // as the span was closed.
        assert!(trace_ctx.host_spans.len() == 1);
    }

    /// Test handling a batch with one span and one event.
    /// The span is not closed.
    #[test]
    fn test_guest_trace_span_with_event() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "test-span", "test-target", 2000, vec![]),
            create_log_event(1, 2500, "test-event", vec![]),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.len() == 1);
        // The active host span is new because a new guest span was created
        assert!(trace_ctx.host_spans.len() == 2);
    }

    /// Test handling a batch with two open spans in a parent-child relationship.
    /// The spans are not closed.
    #[test]
    fn test_guest_trace_parent_child_spans() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "parent-span", "test-target", 2000, vec![]),
            create_open_span(2, Some(1), "child-span", "test-target", 2500, vec![]),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.len() == 2);
        // The active host span is new because new guest spans were created
        assert!(trace_ctx.host_spans.len() == 2);
    }

    /// Test handling a batch with two closed spans in a parent-child relationship.
    /// The spans are closed.
    #[test]
    fn test_guest_trace_closed_parent_child_spans() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "parent-span", "test-target", 2000, vec![]),
            create_open_span(2, Some(1), "child-span", "test-target", 2500, vec![]),
            create_close_span(2, 3000),
            create_close_span(1, 3500),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.is_empty());
        // The active host span is the same as before because no new guest spans were created
        // as the spans were closed.
        assert!(trace_ctx.host_spans.len() == 1);
    }

    /// Test handling a batch with two spans partially closed in a parent-child
    /// relationship.
    /// The parent span is open, the child span is closed.
    #[test]
    fn test_guest_trace_partially_closed_parent_child_spans() {
        let mut trace_ctx = create_dummy_trace_context();

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "parent-span", "test-target", 2000, vec![]),
            create_open_span(2, Some(1), "child-span", "test-target", 2500, vec![]),
            create_close_span(2, 3000),
        ];

        let res = trace_ctx.handle_trace_impl(events);
        assert!(res.is_ok());
        assert!(trace_ctx.guest_spans.len() == 1);
        // The active host span is new because a new guest span was created
        assert!(trace_ctx.host_spans.len() == 2);
    }

    #[test]
    fn test_guest_trace_span_without_guest_start_errors() {
        let mut trace_ctx = TraceContext::new();
        trace_ctx.tsc_freq = Some(3_200_000_000);
        trace_ctx.start_wall = Some(SystemTime::now());
        trace_ctx.start_instant = Some(Instant::now());

        let events = vec![create_open_span(1, None, "span", "target", 2000, vec![])];

        let err = trace_ctx
            .handle_trace_impl(events)
            .expect_err("Span before GuestStart must error");
        assert!(
            err.to_string()
                .contains("Guest start TSC not set before opening guest span"),
            "unexpected error: {}",
            err,
        );
        assert!(trace_ctx.guest_spans.is_empty());
        assert_eq!(trace_ctx.host_spans.len(), 1);
    }

    #[test]
    fn test_guest_trace_missing_start_wall_errors() {
        let mut trace_ctx = TraceContext::new();
        trace_ctx.tsc_freq = Some(3_200_000_000);
        trace_ctx.start_tsc = Some(1000);
        trace_ctx.start_instant = Some(Instant::now());

        let events = vec![
            GuestEvent::GuestStart { tsc: 1000 },
            create_open_span(1, None, "span", "target", 1500, vec![]),
        ];

        let err = trace_ctx
            .handle_trace_impl(events)
            .expect_err("Missing start_wall should error");
        assert!(
            err.to_string().contains("start_wall not set"),
            "unexpected error: {}",
            err,
        );
    }

    #[test]
    fn test_calculate_guest_time_requires_tsc_freq() {
        let mut trace_ctx = TraceContext::new();
        trace_ctx.start_wall = Some(SystemTime::now());

        let err = trace_ctx
            .calculate_guest_time_relative_to_host(0, 0)
            .expect_err("Missing TSC frequency should error");
        assert!(
            err.to_string().contains("TSC frequency not set"),
            "unexpected error: {}",
            err,
        );
    }

    #[test]
    fn test_calculate_guest_time_requires_start_wall() {
        let mut trace_ctx = TraceContext::new();
        trace_ctx.tsc_freq = Some(3_200_000_000);

        let err = trace_ctx
            .calculate_guest_time_relative_to_host(0, 0)
            .expect_err("Missing start wall time should error");
        assert!(
            err.to_string().contains("start_wall not set"),
            "unexpected error: {}",
            err,
        );
    }
}