supermachine 0.7.108

//! Full stopped-snapshot restore orchestration for the kboxlike backend.
//!
//! Lower layers own the details: fd snapshots rebuild host shadows, tracee
//! restore installs memory/kernel/register state, and the shadow rebuilder
//! replays kernel-side fd relationships. This module keeps the production
//! sequence in one place.

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::fd::{
    FdError, KboxlikeFdSnapshot, KboxlikeFdSystem, ModelThreadId, ShadowBacking,
    ShadowInstallPlanWithBacking,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::shadow::{HostShadowRebuilder, KboxlikeFdKernelReplaySummary};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::snapshot::{
    KboxlikeEpollShadowBackingPlan, KboxlikeFdQueuedDataPlan, KboxlikeFdScmRightsPlan,
    KboxlikeFdSocketOptionsPlan, KboxlikeFdSocketShutdownPlan, KboxlikeMaterializedMemoryRestore,
    KboxlikeSharedMemoryDump, KboxlikeThreadKernelStateSnapshot, KboxlikeThreadSnapshot,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::tracee::{
    apply_shadow_install_plan_with_scratch,
    apply_stopped_tracee_full_restore_set_with_replacement_pids,
    apply_supervisor_fd_scm_rights_install_with_scratch, apply_thread_register_restore_for_tracee,
    apply_tracee_credentials_restore_with_replacement_pids, apply_tracee_mm_exe_file_with_scratch,
    create_stopped_replacement_tracee_set, resume_restored_tracee_set, RestoredTraceeResumer,
    StoppedTraceeReplacementFactory, TraceeFullRestoreSetSummary,
    TraceeReplacementThreadGroupCreateSummary, TraceeRestoreAccessFactory,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use std::collections::{BTreeMap, BTreeSet};

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) struct KboxlikeSnapshotRestoreSummary {
    pub(crate) replacements: TraceeReplacementThreadGroupCreateSummary,
    pub(crate) tracee: TraceeFullRestoreSetSummary,
    pub(crate) fd_tracee_installs: usize,
    pub(crate) fd_kernel: KboxlikeFdKernelReplaySummary,
    pub(crate) resume: super::tracee::TraceeResumeSetSummary,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Debug)]
pub(crate) struct KboxlikeSnapshotRestoreResult {
    pub(crate) fd_system: KboxlikeFdSystem,
    pub(crate) replacement_host_pids: BTreeMap<ModelThreadId, i32>,
    pub(crate) summary: KboxlikeSnapshotRestoreSummary,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Clone, Copy, Debug)]
pub(crate) struct KboxlikeFdKernelReplayInputs<'a> {
    pub(crate) epoll_plans: &'a [KboxlikeEpollShadowBackingPlan],
    pub(crate) socket_options_plans: &'a [KboxlikeFdSocketOptionsPlan],
    pub(crate) queued_data_plans: &'a [KboxlikeFdQueuedDataPlan],
    pub(crate) socket_shutdown_plans: &'a [KboxlikeFdSocketShutdownPlan],
    pub(crate) scm_rights_plans: &'a [KboxlikeFdScmRightsPlan],
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
impl<'a> KboxlikeFdKernelReplayInputs<'a> {
    pub(crate) fn empty() -> Self {
        Self {
            epoll_plans: &[],
            socket_options_plans: &[],
            queued_data_plans: &[],
            socket_shutdown_plans: &[],
            scm_rights_plans: &[],
        }
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
pub(crate) struct KboxlikeTraceeRestoreInputs<'a> {
    pub(crate) memory_restore: &'a KboxlikeMaterializedMemoryRestore,
    pub(crate) threads: &'a [KboxlikeThreadSnapshot],
    pub(crate) kernel_states: &'a [KboxlikeThreadKernelStateSnapshot],
    pub(crate) shared_dumps: &'a [KboxlikeSharedMemoryDump],
    pub(crate) supervisor_pid: i32,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
pub(crate) fn restore_stopped_kboxlike_snapshot_with_replacements(
    fd_snapshot: KboxlikeFdSnapshot,
    tracee: KboxlikeTraceeRestoreInputs<'_>,
    fd_replay: KboxlikeFdKernelReplayInputs<'_>,
    rebuilder: &mut HostShadowRebuilder,
    replacement_factory: &mut impl StoppedTraceeReplacementFactory,
    access_factory: &mut impl TraceeRestoreAccessFactory,
    resumer: &mut impl RestoredTraceeResumer,
) -> Result<KboxlikeSnapshotRestoreResult, FdError> {
    let created = create_stopped_replacement_tracee_set(tracee.threads, replacement_factory)?;
    let fd_system = KboxlikeFdSystem::restore_rebuilding_shadows(fd_snapshot, rebuilder)?;
    let tracee_summary = apply_stopped_tracee_full_restore_set_with_replacement_pids(
        tracee.memory_restore,
        tracee.threads,
        tracee.kernel_states,
        tracee.shared_dumps,
        &created.replacement_host_pids,
        tracee.supervisor_pid,
        access_factory,
    )?;
    let fd_tracee_installs = install_rebuilt_shadow_fds_in_replacements(
        &fd_system,
        tracee.threads,
        &created.replacement_host_pids,
        tracee.supervisor_pid,
        access_factory,
    )?;
    let fd_kernel_summary = rebuilder.replay_fd_kernel_state_with_replacement_pids(
        &fd_system,
        fd_replay.epoll_plans,
        fd_replay.socket_options_plans,
        fd_replay.queued_data_plans,
        fd_replay.socket_shutdown_plans,
        fd_replay.scm_rights_plans,
        &created.replacement_host_pids,
    )?;
    apply_tracee_credentials_restore_with_replacement_pids(
        tracee.threads,
        &created.replacement_host_pids,
        access_factory,
    )?;
    let resume_summary = resume_restored_tracee_set(
        created
            .replacement_host_pids
            .values()
            .copied()
            .collect::<BTreeSet<_>>(),
        resumer,
    )?;

    Ok(KboxlikeSnapshotRestoreResult {
        fd_system,
        replacement_host_pids: created.replacement_host_pids,
        summary: KboxlikeSnapshotRestoreSummary {
            replacements: created.summary,
            tracee: tracee_summary,
            fd_tracee_installs,
            fd_kernel: fd_kernel_summary,
            resume: resume_summary,
        },
    })
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn install_rebuilt_shadow_fds_in_replacements(
    fd_system: &KboxlikeFdSystem,
    threads: &[KboxlikeThreadSnapshot],
    replacement_host_pids: &BTreeMap<ModelThreadId, i32>,
    supervisor_pid: i32,
    access_factory: &mut impl TraceeRestoreAccessFactory,
) -> Result<usize, FdError> {
    let mut installed = 0;
    let mut seen_model_pids = BTreeSet::new();
    for thread in threads {
        if !seen_model_pids.insert(thread.model_pid) {
            continue;
        }
        let Some(replacement_host_pid) =
            replacement_host_pids.get(&thread.model_thread_id).copied()
        else {
            return Err(FdError::TraceeInstallFailed(
                "missing replacement tracee pid for fd install",
            ));
        };
        let plans = fd_system.process_shadow_install_plans_with_backing(thread.model_pid)?;
        if plans.is_empty() {
            continue;
        }
        let mut access = access_factory.access_for_restore(replacement_host_pid)?;
        let plan_sets = partition_scm_install_plans(plans)?;
        for plan in plan_sets.proc_fd {
            apply_shadow_install_plan_with_scratch(
                supervisor_pid,
                plan.install,
                &mut access.writer,
                &mut access.executor,
            )?;
            installed += 1;
        }
        for plan in plan_sets.socketpair_scm {
            apply_supervisor_fd_scm_rights_install_with_scratch(
                plan.install.shadow.supervisor_fd,
                plan.install,
                &mut access.writer,
                &mut access.executor,
            )?;
            installed += 1;
        }
        if let Ok(exe_path) = std::fs::read_link(format!("/proc/{}/exe", thread.host_pid)) {
            apply_tracee_mm_exe_file_with_scratch(
                &exe_path,
                &mut access.writer,
                &mut access.executor,
            )?;
            if std::env::var_os("SUPERMACHINE_KBOXLIKE_TRACE_FD_RESTORE").is_some() {
                eprintln!(
                    "KBOXLIVE_EXE_FILE_RESTORE source_pid={} replacement_pid={} path={}",
                    thread.host_pid,
                    replacement_host_pid,
                    exe_path.display(),
                );
            }
        }
        let restored_thread = KboxlikeThreadSnapshot {
            host_pid: replacement_host_pid,
            model_pid: thread.model_pid,
            model_thread_id: thread.model_thread_id,
            regs: thread.regs,
        };
        apply_thread_register_restore_for_tracee(
            std::slice::from_ref(&restored_thread),
            replacement_host_pid,
            &mut access.register_writer,
        )?;
    }
    Ok(installed)
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Default)]
struct PartitionedInstallPlans {
    proc_fd: Vec<ShadowInstallPlanWithBacking>,
    socketpair_scm: Vec<ShadowInstallPlanWithBacking>,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn partition_scm_install_plans(
    plans: Vec<ShadowInstallPlanWithBacking>,
) -> Result<PartitionedInstallPlans, FdError> {
    let mut partitioned = PartitionedInstallPlans::default();

    for plan in plans {
        if requires_scm_rights_install(plan.backing.as_ref()) {
            partitioned.socketpair_scm.push(plan);
        } else {
            partitioned.proc_fd.push(plan);
        }
    }

    Ok(partitioned)
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn requires_scm_rights_install(backing: Option<&ShadowBacking>) -> bool {
    let Some(ShadowBacking::HostPassthrough { label }) = backing else {
        return false;
    };
    // socketpair/livefd: sockets can never be reopened via /proc/<pid>/fd.
    // eventfd/epoll/timerfd/signalfd/inotify: anon-inode fds — the
    // /proc/<pid>/fd reopen fails with ENXIO on modern kernels, and epoll
    // additionally must be the SAME kernel object the supervisor replays
    // interests into, which only an SCM_RIGHTS dup (not a reopen) preserves.
    label.starts_with("socketpair:")
        || label.starts_with("livefd:")
        || label.starts_with("eventfd:")
        || label.starts_with("epoll:")
        || label.starts_with("timerfd:")
        || label.starts_with("signalfd:")
        || label.starts_with("inotify:")
}

#[cfg(test)]
mod tests {
    #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
    mod linux_x86_64 {
        use super::super::*;
        use crate::kboxlike::fd::{KboxlikeFdSystem, ProcessId};
        use crate::kboxlike::snapshot::{
            KboxlikeByteBackedRestoreSource, KboxlikeMaterializedMemoryRestoreAction,
        };
        use crate::kboxlike::tracee::{
            TraceeKernelStateRestoreSummary, TraceeMemoryRestoreSummary, TraceeMemoryWriter,
            TraceeRegisterRestoreSummary, TraceeRegisterWriter, TraceeRestoreAccess,
            TraceeRestoreAccessFactory, TraceeResumeSetSummary, TraceeSyscall,
            TraceeSyscallExecutor,
        };
        use std::collections::{BTreeMap, VecDeque};

        #[derive(Default)]
        struct FakeMemory {
            writes: Vec<(u64, Vec<u8>)>,
        }

        struct FakeExecutor {
            returns: VecDeque<i64>,
        }

        #[derive(Default)]
        struct FakeRegisterWriter {
            writes: Vec<libc::user_regs_struct>,
        }

        #[derive(Default)]
        struct FakeAccessFactory {
            returns_by_pid: BTreeMap<i32, Vec<i64>>,
            requested: Vec<i32>,
        }

        #[derive(Default)]
        struct FakeReplacementFactory {
            calls: Vec<(ProcessId, Vec<ModelThreadId>)>,
            returns: VecDeque<Result<BTreeMap<ModelThreadId, i32>, FdError>>,
        }

        #[derive(Default)]
        struct FakeResumer {
            resumed: Vec<i32>,
        }

        impl TraceeSyscallExecutor for FakeExecutor {
            fn syscall(&mut self, _syscall: TraceeSyscall) -> Result<i64, FdError> {
                Ok(self.returns.pop_front().unwrap_or(0))
            }
        }

        impl TraceeMemoryWriter for FakeMemory {
            fn write_bytes(&mut self, addr: u64, bytes: &[u8]) -> Result<(), FdError> {
                self.writes.push((addr, bytes.to_vec()));
                Ok(())
            }
        }

        impl TraceeRegisterWriter for FakeRegisterWriter {
            fn write_regs(&mut self, regs: &libc::user_regs_struct) -> Result<(), FdError> {
                self.writes.push(*regs);
                Ok(())
            }
        }

        impl TraceeRestoreAccessFactory for FakeAccessFactory {
            type Writer = FakeMemory;
            type Executor = FakeExecutor;
            type RegisterWriter = FakeRegisterWriter;

            fn access_for_restore(
                &mut self,
                host_pid: i32,
            ) -> Result<
                TraceeRestoreAccess<Self::Writer, Self::Executor, Self::RegisterWriter>,
                FdError,
            > {
                self.requested.push(host_pid);
                let returns = self
                    .returns_by_pid
                    .remove(&host_pid)
                    .ok_or(FdError::TraceeInstallFailed("missing fake tracee returns"))?;
                Ok(TraceeRestoreAccess {
                    writer: FakeMemory::default(),
                    executor: FakeExecutor {
                        returns: returns.into(),
                    },
                    register_writer: FakeRegisterWriter::default(),
                })
            }
        }

        impl StoppedTraceeReplacementFactory for FakeReplacementFactory {
            fn create_stopped_replacement_group(
                &mut self,
                model_pid: ProcessId,
                threads: &[KboxlikeThreadSnapshot],
            ) -> Result<BTreeMap<ModelThreadId, i32>, FdError> {
                self.calls.push((
                    model_pid,
                    threads
                        .iter()
                        .map(|thread| thread.model_thread_id)
                        .collect(),
                ));
                self.returns.pop_front().unwrap_or_else(|| {
                    Err(FdError::TraceeInstallFailed("missing fake replacement"))
                })
            }
        }

        impl RestoredTraceeResumer for FakeResumer {
            fn resume_restored_tracee(&mut self, host_pid: i32) -> Result<(), FdError> {
                self.resumed.push(host_pid);
                Ok(())
            }
        }

        #[test]
        fn snapshot_restore_rebuilds_tracees_fds_kernel_state_then_resumes() {
            let mut fd_system = KboxlikeFdSystem::new();
            let model_pid = fd_system.spawn_init_with_exe("/ms-playwright/chromium/chrome");
            let fd_snapshot = fd_system.snapshot();

            let mut regs: libc::user_regs_struct = unsafe { std::mem::zeroed() };
            regs.rip = 0x7f00_3000;
            // Credential replay reads the source tracee's /proc/<pid>/status.
            // Use this live test process as the credential source instead of a
            // fake pid so the unit test exercises the real /proc parser without
            // depending on an unrelated process existing on the host.
            let source_host_pid = std::process::id() as i32;
            let threads = vec![KboxlikeThreadSnapshot {
                host_pid: source_host_pid,
                model_pid,
                model_thread_id: 1,
                regs,
            }];
            let memory_restore = KboxlikeMaterializedMemoryRestore {
                actions: vec![KboxlikeMaterializedMemoryRestoreAction::ByteBacked {
                    host_pid: source_host_pid,
                    model_pid,
                    start: 0x7f00_3000,
                    prot: libc::PROT_READ | libc::PROT_WRITE,
                    bytes: b"code".to_vec(),
                    source: KboxlikeByteBackedRestoreSource::AnonymousPrivate { label: None },
                }],
            };
            let kernel_states = vec![KboxlikeThreadKernelStateSnapshot {
                host_pid: source_host_pid,
                model_pid,
                model_thread_id: 1,
                clear_child_tid: Some(0x7f00_4000),
                robust_list_head: Some(0x7f00_5000),
                robust_list_len: 24,
                rseq: None,
            }];
            let mut rebuilder = HostShadowRebuilder::new(".");
            let mut replacement_factory = FakeReplacementFactory {
                returns: VecDeque::from([Ok(BTreeMap::from([(1, 9001)]))]),
                ..FakeReplacementFactory::default()
            };
            let mut access_factory = FakeAccessFactory::default();
            access_factory
                .returns_by_pid
                .insert(9001, vec![0x7f00_3000, 9001, 0, 0, 0, 0, 0, 0, 0, 0]);
            let mut resumer = FakeResumer::default();

            let result = restore_stopped_kboxlike_snapshot_with_replacements(
                fd_snapshot,
                KboxlikeTraceeRestoreInputs {
                    memory_restore: &memory_restore,
                    threads: &threads,
                    kernel_states: &kernel_states,
                    shared_dumps: &[],
                    supervisor_pid: 1234,
                },
                KboxlikeFdKernelReplayInputs::empty(),
                &mut rebuilder,
                &mut replacement_factory,
                &mut access_factory,
                &mut resumer,
            )
            .unwrap();

            assert_eq!(replacement_factory.calls, vec![(model_pid, vec![1])]);
            assert_eq!(access_factory.requested, vec![9001]);
            assert_eq!(resumer.resumed, vec![9001]);
            assert_eq!(result.replacement_host_pids, BTreeMap::from([(1, 9001)]));
            assert_eq!(result.fd_system.process_count(), 1);
            assert_eq!(
                result.summary,
                KboxlikeSnapshotRestoreSummary {
                    replacements: TraceeReplacementThreadGroupCreateSummary {
                        process_groups: 1,
                        tracees_requested: 1,
                        tracees_created: 1,
                    },
                    tracee: TraceeFullRestoreSetSummary {
                        tracees_restored: 1,
                        shared_objects: Default::default(),
                        memory: TraceeMemoryRestoreSummary {
                            byte_backed_applied: 1,
                            clean_file_remapped: 0,
                            shared_object_remapped: 0,
                            skipped_no_mapping: 0,
                            deferred_unsupported: 0,
                            other_tracee: 0,
                        },
                        kernel_state: TraceeKernelStateRestoreSummary {
                            clear_child_tid_applied: 1,
                            robust_list_applied: 1,
                            rseq_applied: 0,
                            other_tracee: 0,
                        },
                        registers: TraceeRegisterRestoreSummary {
                            registers_restored: 1,
                            other_tracee: 0,
                        },
                    },
                    fd_tracee_installs: 0,
                    fd_kernel: KboxlikeFdKernelReplaySummary::default(),
                    resume: TraceeResumeSetSummary { tracees_resumed: 1 },
                }
            );
        }
    }
}