lazaret/crates/lazaret-host-coordinator/src/control/execution.rs

use super::{
    Arc, AsFd, CONTROL_SOCKET_NAME, ContainerId, ControlError, File, FileIdentity,
    InteractiveSessionBatch, InteractiveSessionEvent, InteractiveSessionRequest, MetadataExt, Mode,
    NonInteractiveSessionMessage, NonInteractiveSessionRequest, OFlags, OpenOptions, OsStr,
    OsString, OsStringExt, PUBLICATION_IDENTITY_NAME, Path, PathBuf, ProcessExit,
    ProtocolPublicationIdentity, Read, RunscCommand, STAGED_CONTROL_SOCKET_NAME,
    STAGED_PUBLICATION_IDENTITY_NAME, ServerConfig, SessionExit, SessionId, SessionRegistry,
    TargetCommand, TrustAnchor, Write, child_broker, fchmod, fmt, io, openat,
    remove_if_identity_matches,
};

pub(super) struct ControlPaths {
    pub(super) root: PathBuf,
    pub(super) root_file: File,
    pub(super) root_identity: FileIdentity,
    pub(super) socket: PathBuf,
    pub(super) identity: PathBuf,
    pub(super) staged_socket: PathBuf,
    pub(super) staged_identity: PathBuf,
}

impl ControlPaths {
    pub(super) fn open(root: &Path, uid: u32) -> Result<Self, ControlError> {
        if !root.is_absolute() {
            return Err(ControlError::InvalidPath(
                "runtime Domain root is not absolute",
            ));
        }
        let canonical = std::fs::canonicalize(root)?;
        if canonical != root {
            return Err(ControlError::InvalidPath(
                "runtime Domain root contains a symlink",
            ));
        }
        let metadata = std::fs::symlink_metadata(root)?;
        if !metadata.file_type().is_dir()
            || metadata.uid() != uid
            || metadata.mode() & 0o7777 != 0o700
        {
            return Err(ControlError::InvalidPath(
                "runtime Domain root is not invoking-user-owned mode 0700",
            ));
        }
        let root_file = OpenOptions::new().read(true).open(root)?;
        let root_identity = FileIdentity::from_metadata(&root_file.metadata()?);
        Ok(Self {
            root: root.to_path_buf(),
            root_file,
            root_identity,
            socket: root.join(CONTROL_SOCKET_NAME),
            identity: root.join(PUBLICATION_IDENTITY_NAME),
            staged_socket: root.join(STAGED_CONTROL_SOCKET_NAME),
            staged_identity: root.join(STAGED_PUBLICATION_IDENTITY_NAME),
        })
    }

    pub(super) fn validate_root(&self, uid: u32) -> Result<(), ControlError> {
        let descriptor = self.root_file.metadata()?;
        let current = std::fs::symlink_metadata(&self.root)?;
        if FileIdentity::from_metadata(&descriptor) != self.root_identity
            || FileIdentity::from_metadata(&current) != self.root_identity
            || !current.file_type().is_dir()
            || current.uid() != uid
            || current.mode() & 0o7777 != 0o700
        {
            return Err(ControlError::InvalidPath("runtime Domain root changed"));
        }
        Ok(())
    }

    pub(super) fn create_named_identity(
        &self,
        name: &str,
        bytes: &[u8],
        uid: u32,
    ) -> Result<(File, FileIdentity), ControlError> {
        let fd = openat(
            self.root_file.as_fd(),
            name,
            OFlags::CREATE | OFlags::EXCL | OFlags::RDWR | OFlags::CLOEXEC | OFlags::NOFOLLOW,
            Mode::RUSR | Mode::WUSR,
        )
        .map_err(io::Error::from)?;
        let mut file = File::from(fd);
        let identity = FileIdentity::from_metadata(&file.metadata()?);
        let result = (|| {
            fchmod(file.as_fd(), Mode::RUSR | Mode::WUSR).map_err(io::Error::from)?;
            file.write_all(bytes)
                .map_err(|_| ControlError::InvalidPublication("identity write failed"))?;
            file.sync_all()
                .map_err(|_| ControlError::InvalidPublication("identity sync failed"))?;
            self.root_file
                .sync_all()
                .map_err(|_| ControlError::InvalidPublication("directory sync failed"))?;
            self.validate_identity_at(&self.root.join(name), &file, identity, uid, bytes.len())
        })();
        if let Err(error) = result {
            remove_if_identity_matches(&self.root.join(name), identity);
            return Err(error);
        }
        Ok((file, identity))
    }

    pub(super) fn validate_identity(
        &self,
        file: &File,
        expected: FileIdentity,
        uid: u32,
        expected_length: usize,
    ) -> Result<(), ControlError> {
        self.validate_identity_at(&self.identity, file, expected, uid, expected_length)
    }

    pub(super) fn validate_identity_at(
        &self,
        path: &Path,
        file: &File,
        expected: FileIdentity,
        uid: u32,
        expected_length: usize,
    ) -> Result<(), ControlError> {
        self.validate_root(uid)?;
        let descriptor = file.metadata()?;
        let current = std::fs::symlink_metadata(path)?;
        if FileIdentity::from_metadata(&descriptor) != expected
            || FileIdentity::from_metadata(&current) != expected
            || !descriptor.file_type().is_file()
            || !current.file_type().is_file()
            || descriptor.uid() != uid
            || current.uid() != uid
            || descriptor.mode() & 0o7777 != 0o600
            || current.mode() & 0o7777 != 0o600
            || usize::try_from(descriptor.len()).ok() != Some(expected_length)
            || current.len() != descriptor.len()
        {
            return Err(ControlError::InvalidPublication(
                "publication identity changed type, ownership, mode, or content length",
            ));
        }
        Ok(())
    }
}

pub struct PublishedControlFace {
    pub(super) listener: std::os::unix::net::UnixListener,
    pub(super) paths: ControlPaths,
    pub(super) socket_identity: FileIdentity,
    pub(super) identity_file: File,
    pub(super) identity_file_identity: FileIdentity,
    pub(super) publication_identity: ProtocolPublicationIdentity,
    pub(super) publication_bytes: Vec<u8>,
    pub(super) uid: u32,
    pub(super) anchor: TrustAnchor,
    pub(super) server_config: Arc<ServerConfig>,
    pub(super) registry: Arc<SessionRegistry>,
    pub(super) executor: Arc<dyn NonInteractiveSessionExecutor>,
    pub(super) interactive_executor: Arc<dyn InteractiveSessionExecutor>,
}

pub(crate) struct AcceptedControlConnection {
    pub(super) socket: std::os::unix::net::UnixStream,
}

impl AcceptedControlConnection {
    pub(crate) fn cancellation_handle(&self) -> std::io::Result<std::os::unix::net::UnixStream> {
        self.socket.try_clone()
    }
}

pub trait NonInteractiveSessionExecutor: Send + Sync {
    /// # Errors
    /// Returns a runtime, stream, or exact-session execution failure.
    fn execute(
        &self,
        request: &NonInteractiveSessionRequest,
        emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<SessionExit, ControlError>;
}

pub(super) struct ExecutionUnavailable;

impl NonInteractiveSessionExecutor for ExecutionUnavailable {
    fn execute(
        &self,
        _request: &NonInteractiveSessionRequest,
        _emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<SessionExit, ControlError> {
        Err(ControlError::InvalidRequest)
    }
}

pub trait InteractiveControlSession: Send {
    /// # Errors
    /// Returns an exact-session terminal, prompt-routing, or transport failure.
    fn exchange(
        &mut self,
        event: InteractiveSessionEvent,
    ) -> Result<InteractiveSessionBatch, ControlError>;
}

pub trait InteractiveSessionExecutor: Send + Sync {
    /// # Errors
    /// Returns an acquisition or exact-session identity failure.
    fn start(
        &self,
        request: &InteractiveSessionRequest,
    ) -> Result<Box<dyn InteractiveControlSession>, ControlError>;
}

pub(super) struct InteractiveExecutionUnavailable;

impl InteractiveSessionExecutor for InteractiveExecutionUnavailable {
    fn start(
        &self,
        _request: &InteractiveSessionRequest,
    ) -> Result<Box<dyn InteractiveControlSession>, ControlError> {
        Err(ControlError::InvalidRequest)
    }
}

pub struct RunscNonInteractiveExecutor {
    pub(super) program: OsString,
    pub(super) global_arguments: Vec<OsString>,
    pub(super) container: OsString,
    pub(super) runner: Arc<dyn NonInteractiveCommandRunner>,
}

pub(super) trait NonInteractiveCommandRunner: Send + Sync {
    fn run(
        &self,
        target: &TargetCommand,
        session_id: SessionId,
        emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<ProcessExit, ControlError>;
}

struct ChildBrokerCommandRunner;

impl NonInteractiveCommandRunner for ChildBrokerCommandRunner {
    fn run(
        &self,
        target: &TargetCommand,
        session_id: SessionId,
        emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<ProcessExit, ControlError> {
        let mut child = child_broker().run_transient_captured(target)?;
        let mut stdout_done = false;
        let mut stderr_done = false;
        while !(stdout_done && stderr_done && child.process().has_exited()?) {
            let (stdout, stderr) = child.streams();
            stdout_done |= emit_available(stdout, session_id, true, emit)?;
            stderr_done |= emit_available(stderr, session_id, false, emit)?;
            if !(stdout_done && stderr_done) {
                std::thread::yield_now();
            }
        }
        child_broker()
            .reap_status(child.into_process())
            .map_err(Into::into)
    }
}

impl RunscNonInteractiveExecutor {
    /// Derives the exact `runsc exec` prefix from the generation's attached `runsc run` command.
    ///
    /// # Errors
    /// Refuses a non-runsc command or a command without one exact run/container suffix.
    pub fn from_run_command(
        command: &RunscCommand,
        container_id: ContainerId,
    ) -> Result<Self, ControlError> {
        if command.program() != OsStr::new("runsc") {
            return Err(ControlError::InvalidRequest);
        }
        let run = command
            .arguments()
            .iter()
            .position(|argument| argument == OsStr::new("run"))
            .ok_or(ControlError::InvalidRequest)?;
        if run + 3 != command.arguments().len()
            || !command.arguments()[run + 1]
                .as_encoded_bytes()
                .starts_with(b"-bundle=")
        {
            return Err(ControlError::InvalidRequest);
        }
        let container = OsString::from(hex_lower(container_id.as_bytes()));
        if command.arguments()[run + 2] != container {
            return Err(ControlError::InvalidRequest);
        }
        Ok(Self {
            program: command.program().to_owned(),
            global_arguments: command.arguments()[..run].to_vec(),
            container,
            runner: Arc::new(ChildBrokerCommandRunner),
        })
    }
}

impl NonInteractiveSessionExecutor for RunscNonInteractiveExecutor {
    fn execute(
        &self,
        request: &NonInteractiveSessionRequest,
        emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
    ) -> Result<SessionExit, ControlError> {
        let mut target = TargetCommand::new(self.program.clone());
        for argument in &self.global_arguments {
            target = target.argument(argument.clone());
        }
        target = target
            .argument("exec")
            .argument("-env")
            .argument(format!(
                "LAZARET_SESSION_ID={}",
                hex_lower(request.session_id.as_bytes())
            ))
            .argument(self.container.clone());
        for argument in &request.argv {
            target = target.argument(OsString::from_vec(argument.clone()));
        }
        let process_status = self.runner.run(&target, request.session_id, emit)?;
        match (process_status.exit_code, process_status.terminating_signal) {
            (Some(code), None) => Ok(SessionExit::exited(request.session_id, code)),
            (None, Some(_)) => Ok(SessionExit::terminated(request.session_id)),
            _ => Err(ControlError::InvalidRequest),
        }
    }
}

fn emit_available(
    stream: &mut File,
    session_id: SessionId,
    stdout: bool,
    emit: &mut dyn FnMut(NonInteractiveSessionMessage) -> Result<(), ControlError>,
) -> Result<bool, ControlError> {
    let mut bytes = [0_u8; 16 * 1024];
    loop {
        match stream.read(&mut bytes) {
            Ok(0) => return Ok(true),
            Ok(length) => {
                let output = lazaret_core::records::NonInteractiveOutput {
                    session_id,
                    bytes: bytes[..length].to_vec(),
                };
                emit(if stdout {
                    NonInteractiveSessionMessage::Stdout(output)
                } else {
                    NonInteractiveSessionMessage::Stderr(output)
                })?;
            }
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => return Ok(false),
            Err(error) => return Err(error.into()),
        }
    }
}

pub(super) fn hex_lower(bytes: &[u8]) -> String {
    use fmt::Write as _;

    bytes.iter().fold(String::new(), |mut output, byte| {
        write!(output, "{byte:02x}").expect("writing to a String cannot fail");
        output
    })
}