microsandbox-runtime 0.7.9

//! Ephemeral authority for signals that bypass a backpressured agent socket.

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;

use bytes::Bytes;
use microsandbox_protocol::exec::ExecSignal;
use microsandbox_protocol::exec_control::{
    EXEC_CONTROL_VERSION, ExecControlRequest, ExecControlResponse,
};
use microsandbox_protocol::{
    codec,
    message::{Message, MessageType},
};
use rand::RngExt as _;
use tokio::sync::{Notify, oneshot};

use super::finished_ids::FinishedIds;
use super::relay::{ControlWrite, ControlWriter};

//--------------------------------------------------------------------------------------------------
// Constants
//--------------------------------------------------------------------------------------------------

const DELIVERY_DEADLINE: Duration = Duration::from_secs(3);

//--------------------------------------------------------------------------------------------------
// Types
//--------------------------------------------------------------------------------------------------

/// One relay lifetime; tokens are never reused across restart or restore.
#[derive(Default)]
pub(crate) struct ExecControlRegistry {
    clients: Mutex<HashMap<[u8; 16], Arc<Owner>>>,
    /// Enabled only after the guest acknowledges its bounded input and range-lease contracts.
    pub(crate) enabled: AtomicBool,
    pub(crate) stopping: AtomicBool,
}

struct Owner {
    start: u32,
    end: u32,
    live: AtomicBool,
    changed: Notify,
    executions: Mutex<Executions>,
}

struct Executions {
    active: HashMap<u32, Arc<ExecControlLease>>,
    finished: FinishedIds,
}

/// Dropping the primary connection revokes its sideband authority before slot reuse.
pub(crate) struct ExecControlConnection {
    token: [u8; 16],
    owner: Arc<Owner>,
    registry: Weak<ExecControlRegistry>,
}

/// Attached to a queued signal so retirement also invalidates previously admitted requests.
pub(crate) struct ExecControlLease {
    live: AtomicBool,
    started: AtomicBool,
}

//--------------------------------------------------------------------------------------------------
// Methods
//--------------------------------------------------------------------------------------------------

impl ExecControlRegistry {
    pub(crate) fn connect(self: &Arc<Self>, start: u32, end: u32) -> ExecControlConnection {
        let owner = Arc::new(Owner {
            start,
            end,
            live: AtomicBool::new(true),
            changed: Notify::new(),
            executions: Mutex::new(Executions {
                active: HashMap::new(),
                finished: FinishedIds::new(start..end),
            }),
        });
        let mut clients = self.clients.lock().unwrap();
        let token = loop {
            let token = rand::rng().random::<[u8; 16]>();
            if !clients.contains_key(&token) {
                break token;
            }
        };
        clients.insert(token, Arc::clone(&owner));
        ExecControlConnection {
            token,
            owner,
            registry: Arc::downgrade(self),
        }
    }

    pub(crate) async fn signal(
        &self,
        writer: &ControlWriter,
        request: ExecControlRequest,
    ) -> ExecControlResponse {
        if request.version != EXEC_CONTROL_VERSION {
            return ExecControlResponse::error(
                "unsupported_feature",
                "unsupported exec control contract",
            );
        }
        // Ordinary exec historically accepted zero. This preserves that input without
        // promising a liveness query: delivery still does not acknowledge guest handling.
        if !(0..=64).contains(&request.signal) {
            return ExecControlResponse::error("invalid_signal", "signal must be between 0 and 64");
        }
        let owner = {
            let clients = self.clients.lock().unwrap();
            let Some(owner) = clients.get(&request.connection) else {
                return ExecControlResponse::error(
                    "execution_closed",
                    "the owning agent connection has ended",
                );
            };
            if !(owner.start..owner.end).contains(&request.id) {
                return ExecControlResponse::error(
                    "execution_closed",
                    "execution does not belong to this connection",
                );
            }
            Arc::clone(owner)
        };

        let deadline = tokio::time::Instant::now() + DELIVERY_DEADLINE;
        let enqueue = async {
            // The separate sockets can race, including an immediate kill after exec() returns.
            // Wait for this owner's actual ExecStarted rather than overtaking process creation.
            // Once captured, retain the same lease through terminal/disconnect; never retarget it.
            let mut lease = None;
            let lease = loop {
                let changed = owner.changed.notified();
                tokio::pin!(changed);
                changed.as_mut().enable();

                if !owner.live.load(Ordering::Acquire) {
                    return Err(ExecControlResponse::error(
                        "execution_closed",
                        "the owning agent connection has ended",
                    ));
                }

                if lease.is_none() {
                    let executions = owner.executions.lock().unwrap();
                    if executions.is_finished(request.id) {
                        return Err(ExecControlResponse::error(
                            "execution_closed",
                            "execution has ended",
                        ));
                    }
                    lease = executions.active.get(&request.id).cloned();
                }

                if let Some(lease) = lease.as_ref() {
                    if !lease.live() {
                        return Err(ExecControlResponse::error(
                            "execution_closed",
                            "execution has ended",
                        ));
                    }
                    if lease.started.load(Ordering::Acquire) {
                        break Arc::clone(lease);
                    }
                }
                changed.await;
            };

            let message = Message::with_payload(
                MessageType::ExecSignal,
                request.id,
                &ExecSignal {
                    signal: request.signal,
                },
            )
            .map_err(|error| {
                ExecControlResponse::error("delivery_unconfirmed", error.to_string())
            })?;

            let mut data = Vec::new();
            codec::encode_to_buf(&message, &mut data).map_err(|error| {
                ExecControlResponse::error("delivery_unconfirmed", error.to_string())
            })?;

            let (completion, completed) = oneshot::channel();
            let write = ControlWrite::exec_signal(Bytes::from(data), request.id, lease, completion);
            writer.send(write).await.map_err(|_| {
                ExecControlResponse::error(
                    "delivery_unconfirmed",
                    "exec control transport is closed",
                )
            })?;

            Ok(completed)
        };

        let completed = match tokio::time::timeout_at(deadline, enqueue).await {
            Ok(Ok(completed)) => completed,
            Ok(Err(response)) => return response,
            Err(_) => {
                return ExecControlResponse::error(
                    "delivery_unconfirmed",
                    "signal delivery deadline expired; the signal was not queued and will not be delivered",
                );
            }
        };

        // Enqueueing and delivery confirmation share one deadline.
        match tokio::time::timeout_at(deadline, completed).await {
            Ok(Ok(())) => ExecControlResponse::delivered(),
            Ok(Err(_)) => ExecControlResponse::error(
                "delivery_unconfirmed",
                "execution ended or its transport closed before delivery confirmation",
            ),
            Err(_) => ExecControlResponse::error(
                "delivery_unconfirmed",
                "signal delivery deadline expired; the queued signal may still be delivered",
            ),
        }
    }
}

impl Executions {
    fn is_finished(&self, id: u32) -> bool {
        self.finished.is_finished(id)
    }

    fn mark_finished(&mut self, id: u32) {
        if let Some(lease) = self.active.remove(&id) {
            lease.live.store(false, Ordering::Release);
        }
        // Host rejection can end an execution before it ever gets an active lease.
        self.finished.mark_finished(id);
    }
}

impl ExecControlConnection {
    pub(crate) fn token(&self) -> [u8; 16] {
        self.token
    }

    pub(crate) fn register(&self, id: u32) {
        if !(self.owner.start..self.owner.end).contains(&id) {
            return;
        }
        let mut executions = self.owner.executions.lock().unwrap();
        // Never replace an identity, live or finished, if a malformed client reuses a correlation.
        if executions.is_finished(id) {
            return;
        }

        executions.active.entry(id).or_insert_with(|| {
            Arc::new(ExecControlLease {
                live: AtomicBool::new(true),
                started: AtomicBool::new(false),
            })
        });
        drop(executions);
        self.owner.changed.notify_waiters();
    }

    pub(crate) fn pending(&self, id: u32) -> bool {
        self.owner
            .executions
            .lock()
            .unwrap()
            .active
            .get(&id)
            .is_some_and(|lease| !lease.started.load(Ordering::Acquire))
    }

    pub(crate) fn started(&self, id: u32) {
        if let Some(lease) = self.owner.executions.lock().unwrap().active.get(&id) {
            lease.started.store(true, Ordering::Release);
        }
        self.owner.changed.notify_waiters();
    }

    pub(crate) fn mark_finished(&self, id: u32) {
        self.owner.executions.lock().unwrap().mark_finished(id);
        self.owner.changed.notify_waiters();
    }
}

impl ExecControlLease {
    pub(crate) fn live(&self) -> bool {
        self.live.load(Ordering::Acquire)
    }
}

//--------------------------------------------------------------------------------------------------
// Trait Implementations
//--------------------------------------------------------------------------------------------------

impl Drop for ExecControlConnection {
    fn drop(&mut self) {
        if let Some(registry) = self.registry.upgrade() {
            // Revoke first; an in-flight handler holding a lease is fenced separately below.
            registry.clients.lock().unwrap().remove(&self.token);
        }
        self.owner.live.store(false, Ordering::Release);
        for lease in self.owner.executions.lock().unwrap().active.values() {
            lease.live.store(false, Ordering::Release);
        }
        self.owner.changed.notify_waiters();
    }
}