Skip to main content

mj_controller/
checkpoint_transfer.rs

1//! Controller-side checkpoint transport and verified teardown gates.
2use crate::targets::{
3    CommandExecutor, CommandPlan, CommandSpec, SshTarget, TargetLocator, join_remote_command,
4    worker_root,
5};
6use anyhow::{Context, Result, bail, ensure};
7use mj_checkpoint::archive::validate_component;
8use mj_checkpoint::checkpoint::*;
9use std::fs;
10use std::path::{Path, PathBuf};
11/// Export by streaming the spec to the worker's standard input.
12///
13/// Every wrapper this builds keeps the target's stdin attached: the container
14/// engines are invoked with `exec -i` and `ssh` forwards stdin by default.
15pub fn export_stdin_command(locator: &TargetLocator, session_id: &str) -> Result<CommandSpec> {
16    export_command(locator, session_id, EXPORT_SPEC_STDIN)
17}
18
19pub fn capture_stdin_command(locator: &TargetLocator, session_id: &str) -> Result<CommandSpec> {
20    checkpoint_stdin_command(
21        locator,
22        session_id,
23        "capture-checkpoint",
24        "capture target checkpoint",
25    )
26}
27
28pub fn pack_stdin_command(locator: &TargetLocator, session_id: &str) -> Result<CommandSpec> {
29    checkpoint_stdin_command(
30        locator,
31        session_id,
32        "pack-checkpoint",
33        "pack target checkpoint",
34    )
35}
36
37fn checkpoint_stdin_command(
38    locator: &TargetLocator,
39    session_id: &str,
40    subcommand: &str,
41    purpose: &str,
42) -> Result<CommandSpec> {
43    let root = worker_root(locator, session_id)?;
44    let args = vec![format!("{root}/hel"), "worker".into(), subcommand.into()];
45    crate::targets::command_on_locator(locator, session_id, args, purpose)
46}
47
48pub fn export_command(
49    locator: &TargetLocator,
50    session_id: &str,
51    spec_path: &str,
52) -> Result<CommandSpec> {
53    validate_remote_path(spec_path)?;
54    let root = worker_root(locator, session_id)?;
55    let args = vec![
56        format!("{root}/hel"),
57        "worker".into(),
58        "export-checkpoint".into(),
59        "--spec".into(),
60        spec_path.into(),
61    ];
62    crate::targets::command_on_locator(locator, session_id, args, "export target checkpoint")
63}
64
65pub fn restore_command(
66    locator: &TargetLocator,
67    session_id: &str,
68    spec_path: &str,
69) -> Result<CommandSpec> {
70    validate_remote_path(spec_path)?;
71    let root = worker_root(locator, session_id)?;
72    let args = vec![
73        format!("{root}/hel"),
74        "worker".into(),
75        "restore-checkpoint".into(),
76        "--spec".into(),
77        spec_path.into(),
78    ];
79    crate::targets::command_on_locator(locator, session_id, args, "restore target checkpoint")
80}
81
82#[derive(Debug, Clone)]
83pub struct CheckpointTransfer<'a> {
84    pub locator: &'a TargetLocator,
85    pub session_id: &'a str,
86    pub operation_id: &'a str,
87    pub remote_archive: &'a str,
88    pub destination: &'a Path,
89    pub expected_sha256: &'a str,
90    pub expected_event_frontier: u64,
91    pub expected_event_frontier_digest: &'a str,
92}
93
94/// Unforgeable outside this module: proof that a controller-local archive has
95/// the exact digest reported by the target after its atomic install.
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct VerifiedCheckpoint {
98    session_id: String,
99    archive_path: PathBuf,
100    sha256: String,
101    event_frontier: u64,
102    event_frontier_digest: String,
103}
104
105impl VerifiedCheckpoint {
106    pub fn archive_path(&self) -> &Path {
107        &self.archive_path
108    }
109    pub fn sha256(&self) -> &str {
110        &self.sha256
111    }
112    pub fn event_frontier(&self) -> u64 {
113        self.event_frontier
114    }
115    pub fn event_frontier_digest(&self) -> &str {
116        &self.event_frontier_digest
117    }
118    pub const fn teardown_allowed(&self) -> bool {
119        true
120    }
121}
122
123impl CheckpointTransfer<'_> {
124    pub fn execute(&self, executor: &impl CommandExecutor) -> Result<VerifiedCheckpoint> {
125        validate_remote_path(self.remote_archive)?;
126        let parent = self.destination.parent().unwrap_or_else(|| Path::new("."));
127        fs::create_dir_all(parent)?;
128        let temporary = tempfile::Builder::new()
129            .prefix(".hel-checkpoint-")
130            .tempfile_in(parent)?;
131        let path = temporary.path().to_path_buf();
132        let staging = remote_staging_path(self.session_id, self.operation_id)?;
133        let transfer_result = transfer_plan(
134            self.locator,
135            self.session_id,
136            self.remote_archive,
137            &path,
138            &staging,
139        )?
140        .execute(executor)
141        .context("download target checkpoint");
142        let staging_cleanup_result = cleanup_transfer_staging(self.locator, &staging, executor);
143        if let Err(error) = transfer_result {
144            return match staging_cleanup_result {
145                Ok(()) => Err(error),
146                Err(cleanup) => Err(error.context(format!(
147                    "clean target checkpoint host staging also failed: {cleanup:#}"
148                ))),
149            };
150        }
151        staging_cleanup_result.context("clean target checkpoint host staging")?;
152        let sha256 = checkpoint_sha256(&path).context("hash downloaded checkpoint")?;
153        ensure!(
154            sha256 == self.expected_sha256,
155            "target and controller checkpoint checksums differ for complete checkpoint archive: \
156             session={}, operation={}, expected_sha256={}, downloaded_sha256={}, downloaded_bytes={}; \
157             target archive retained at {}. The previous verified checkpoint was not replaced and \
158             the source workspace was not removed. Preserve the source and retry a fresh export; \
159             if this repeats, inspect the retained archive and its transfer path",
160            self.session_id,
161            self.operation_id,
162            self.expected_sha256,
163            sha256,
164            fs::metadata(&path)
165                .context("stat downloaded checkpoint")?
166                .len(),
167            self.remote_archive,
168        );
169        temporary
170            .persist(self.destination)
171            .map_err(|error| error.error)?;
172        // The bytes were already checksum-verified in this same directory and
173        // the rename is atomic, so installation only has to make the copy
174        // private and durable; re-reading it would hash the same archive again.
175        let post_install = (|| -> Result<()> {
176            restrict_permissions(self.destination)?;
177            sync_directory(parent)
178        })();
179        if let Err(error) = post_install {
180            return Err(remove_failed_checkpoint_install(self.destination, error));
181        }
182        Ok(VerifiedCheckpoint {
183            session_id: self.session_id.to_owned(),
184            archive_path: self.destination.to_path_buf(),
185            sha256,
186            event_frontier: self.expected_event_frontier,
187            event_frontier_digest: self.expected_event_frontier_digest.to_owned(),
188        })
189    }
190
191    pub fn cleanup_plan(&self, gate: &VerifiedCheckpoint) -> Result<CommandPlan> {
192        ensure!(
193            gate.session_id == self.session_id,
194            "checkpoint gate belongs to another session"
195        );
196        cleanup_plan(self.locator, self.session_id, self.remote_archive)
197    }
198}
199
200/// SSH container transfers use a host-side copy because `scp` cannot address
201/// a path inside the container. The copy is disposable and must be removed as
202/// soon as it has been downloaded; the in-container archive remains gated by
203/// [`CheckpointTransfer::cleanup_plan`] until the local copy is verified.
204fn cleanup_transfer_staging(
205    locator: &TargetLocator,
206    staging: &str,
207    executor: &impl CommandExecutor,
208) -> Result<()> {
209    validate_remote_path(staging)?;
210    let command = match locator {
211        TargetLocator::SshPodman { ssh, .. } | TargetLocator::SshDocker { ssh, .. } => Some(
212            ssh_command(ssh, ["rm", "-f", "--", staging])
213                .purpose("remove remote checkpoint staging"),
214        ),
215        _ => None,
216    };
217    if let Some(command) = command {
218        let output = executor.execute(&command)?;
219        if output.status != 0 {
220            bail!(
221                "{} failed with status {}: {}",
222                command.purpose,
223                output.status,
224                String::from_utf8_lossy(&output.stderr)
225            );
226        }
227    }
228    Ok(())
229}
230
231pub fn transfer_plan(
232    locator: &TargetLocator,
233    session_id: &str,
234    remote_archive: &str,
235    local_temporary: &Path,
236    staging: &str,
237) -> Result<CommandPlan> {
238    validate_remote_path(remote_archive)?;
239    validate_remote_path(staging)?;
240    ensure!(
241        local_temporary.is_absolute(),
242        "local temporary path must be absolute"
243    );
244    worker_root(locator, session_id)?;
245    let local = local_temporary.to_string_lossy().into_owned();
246    let mut commands = match locator {
247        TargetLocator::LocalBare { .. } => vec![
248            CommandSpec::new("cp", [remote_archive, local.as_str()])
249                .purpose("copy local bare checkpoint"),
250        ],
251        TargetLocator::LocalPodman { container_id, .. } => vec![
252            CommandSpec::new(
253                "podman",
254                ["cp", &format!("{container_id}:{remote_archive}"), &local],
255            )
256            .purpose("download checkpoint from local Podman"),
257        ],
258        TargetLocator::LocalDocker { container_id } => vec![
259            CommandSpec::new(
260                "docker",
261                ["cp", &format!("{container_id}:{remote_archive}"), &local],
262            )
263            .purpose("download checkpoint from local Docker"),
264        ],
265        TargetLocator::AppleContainer { container_id } => vec![
266            CommandSpec::new(
267                "container",
268                ["cp", &format!("{container_id}:{remote_archive}"), &local],
269            )
270            .purpose("download checkpoint from Apple container"),
271        ],
272        TargetLocator::AwsEc2 { ssh, .. } | TargetLocator::SshBare { ssh, .. } => {
273            vec![scp_command(ssh, remote_archive, &local).purpose("download checkpoint over SSH")]
274        }
275        TargetLocator::SshPodman {
276            ssh, container_id, ..
277        }
278        | TargetLocator::SshDocker { ssh, container_id } => {
279            vec![
280                ssh_command(ssh, ["mkdir", "-p", ".local/share/hel/transfers"])
281                    .purpose("create remote checkpoint staging directory"),
282                ssh_command(
283                    ssh,
284                    [
285                        locator.container_engine().expect("remote container"),
286                        "cp",
287                        &format!("{container_id}:{remote_archive}"),
288                        staging,
289                    ],
290                )
291                .purpose("stage remote container checkpoint"),
292            ]
293        }
294    };
295    if let TargetLocator::SshPodman { ssh, .. } | TargetLocator::SshDocker { ssh, .. } = locator {
296        commands.push(
297            scp_command(ssh, staging, &local)
298                .purpose("download remote container checkpoint over SSH"),
299        );
300    }
301    Ok(CommandPlan {
302        description: format!("download checkpoint for {session_id}"),
303        commands,
304    })
305}
306
307fn cleanup_plan(locator: &TargetLocator, session_id: &str, remote: &str) -> Result<CommandPlan> {
308    validate_remote_path(remote)?;
309    worker_root(locator, session_id)?;
310    let commands = match locator {
311        TargetLocator::LocalBare { .. } => vec![
312            CommandSpec::new("rm", ["-f", "--", remote])
313                .purpose("remove local bare checkpoint staging"),
314        ],
315        TargetLocator::LocalPodman { container_id, .. } => vec![container_exec(
316            "podman",
317            container_id,
318            ["rm", "-f", "--", remote],
319        )],
320        TargetLocator::LocalDocker { container_id } => vec![container_exec(
321            "docker",
322            container_id,
323            ["rm", "-f", "--", remote],
324        )],
325        TargetLocator::AppleContainer { container_id } => vec![container_exec(
326            "container",
327            container_id,
328            ["rm", "-f", "--", remote],
329        )],
330        TargetLocator::AwsEc2 { ssh, .. } | TargetLocator::SshBare { ssh, .. } => {
331            vec![ssh_command(ssh, ["rm", "-f", "--", remote])]
332        }
333        TargetLocator::SshPodman {
334            ssh, container_id, ..
335        }
336        | TargetLocator::SshDocker { ssh, container_id } => vec![ssh_command(
337            ssh,
338            [
339                locator.container_engine().expect("remote container"),
340                "exec",
341                container_id,
342                "rm",
343                "-f",
344                "--",
345                remote,
346            ],
347        )],
348    };
349    Ok(CommandPlan {
350        description: format!("clean checkpoint for {session_id}"),
351        commands,
352    })
353}
354
355fn remote_staging_path(session_id: &str, operation_id: &str) -> Result<String> {
356    validate_component(session_id, "session ID")?;
357    validate_component(operation_id, "checkpoint operation ID")?;
358    Ok(format!(
359        ".local/share/hel/transfers/{session_id}-{operation_id}.hel.zip"
360    ))
361}
362
363fn scp_command(ssh: &SshTarget, remote: &str, local: &str) -> CommandSpec {
364    let mut args = ssh.ssh_args.clone();
365    for argument in &mut args {
366        if argument == "-p" {
367            *argument = "-P".into();
368        }
369    }
370    args.push(format!("{}:{remote}", ssh.destination));
371    args.push(local.into());
372    // `scp` opens its own connection to the same host, so it competes for the
373    // same pre-auth budget and is admitted and retried the same way.
374    CommandSpec::new("scp", args).ssh_destination(ssh.destination.clone())
375}
376
377fn ssh_command(ssh: &SshTarget, args: impl IntoIterator<Item = impl AsRef<str>>) -> CommandSpec {
378    let remote = args
379        .into_iter()
380        .map(|arg| arg.as_ref().to_owned())
381        .collect::<Vec<_>>();
382    let mut command = ssh.ssh_args.clone();
383    command.push(ssh.destination.clone());
384    command.push(join_remote_command(&remote));
385    CommandSpec::new("ssh", command).ssh_destination(ssh.destination.clone())
386}
387
388fn container_exec(
389    engine: &str,
390    id: &str,
391    args: impl IntoIterator<Item = impl Into<String>>,
392) -> CommandSpec {
393    let mut command = vec!["exec".into(), "-i".into(), id.into()];
394    command.extend(args.into_iter().map(Into::into));
395    CommandSpec::new(engine, command)
396}
397
398fn validate_remote_path(path: &str) -> Result<()> {
399    ensure!(!path.is_empty());
400    ensure!(
401        path.bytes()
402            .all(|byte| byte.is_ascii_alphanumeric()
403                || matches!(byte, b'/' | b'~' | b'.' | b'-' | b'_')),
404        "unsafe remote path"
405    );
406    ensure!(
407        !path.split('/').any(|component| component == ".."),
408        "remote path traverses parent"
409    );
410    Ok(())
411}
412
413#[cfg(test)]
414mod tests {
415    use mj_core::config::HarnessKind;
416    use serde_json::json;
417
418    use super::*;
419    use mj_checkpoint::archive::*;
420    use mj_worker::checkpoint::*;
421    use std::cell::RefCell;
422    use std::process::Command;
423
424    use mj_checkpoint::archive::{
425        CanonicalExecutionState, CanonicalQueuedCommandKind, CanonicalQueuedPrompt,
426        CanonicalSessionState, CanonicalTranscriptItem,
427    };
428    use mj_core::targets::CommandOutput;
429
430    const SESSION: &str = "018f9dd2-a3b4-7c8d-9000-123456789abc";
431
432    /// A checkpoint download opens its own connection to the host, so it has
433    /// to be admitted and retried the same way an `ssh` invocation is.
434    #[test]
435    fn a_checkpoint_scp_is_tagged_with_the_connection_destination() {
436        let ssh = SshTarget {
437            destination: "build@10.0.0.1".into(),
438            ssh_args: vec!["-p".into(), "2222".into()],
439        };
440
441        let download = scp_command(&ssh, "remote/archive.zip", "/tmp/local.zip");
442
443        assert_eq!(download.program, "scp");
444        assert_eq!(download.ssh_destination.as_deref(), Some("build@10.0.0.1"));
445    }
446    const NATIVE: &str = "0190aabb-ccdd-7eef-9000-abcdef012345";
447
448    fn ssh() -> SshTarget {
449        SshTarget {
450            destination: "dev@example.test".into(),
451            ssh_args: vec!["-p".into(), "2222".into()],
452        }
453    }
454
455    fn locators() -> Vec<TargetLocator> {
456        let name = mj_core::targets::resource_name(SESSION).unwrap();
457        vec![
458            TargetLocator::LocalBare {
459                worker_root: format!("/var/lib/hel/workers/{SESSION}"),
460            },
461            TargetLocator::LocalPodman {
462                container_id: name.clone(),
463                workspace_storage: Default::default(),
464            },
465            TargetLocator::AppleContainer {
466                container_id: name.clone(),
467            },
468            TargetLocator::AwsEc2 {
469                profile: "default".into(),
470                region: "us-east-1".into(),
471                instance_id: "i-0123456789abcdef0".into(),
472                ssh: ssh(),
473                workspace: format!("~/hel/{SESSION}"),
474            },
475            TargetLocator::SshBare {
476                worker_id: None,
477                ssh: ssh(),
478                workspace: format!("~/hel/{SESSION}"),
479            },
480            TargetLocator::SshPodman {
481                ssh: ssh(),
482                container_id: name,
483                workspace_storage: Default::default(),
484            },
485        ]
486    }
487
488    #[test]
489    fn transfer_plans_cover_all_target_boundaries() {
490        let locators = locators();
491        let plans = locators
492            .iter()
493            .map(|locator| {
494                transfer_plan(
495                    locator,
496                    SESSION,
497                    "/var/lib/hel/workers/checkpoint.hel.zip",
498                    Path::new("/var/tmp/checkpoint.zip"),
499                    &remote_staging_path(SESSION, "test-transfer").unwrap(),
500                )
501                .unwrap()
502            })
503            .collect::<Vec<_>>();
504        assert_eq!(plans[0].commands[0].program, "cp");
505        assert_eq!(plans[1].commands[0].program, "podman");
506        assert_eq!(plans[2].commands[0].program, "container");
507        assert_eq!(plans[3].commands[0].program, "scp");
508        assert_eq!(plans[4].commands[0].program, "scp");
509        assert_eq!(plans[5].commands.len(), 3);
510        assert!(
511            plans[5].commands[1]
512                .args
513                .last()
514                .unwrap()
515                .contains("'podman' 'cp'")
516        );
517        assert!(
518            !plans[5]
519                .commands
520                .iter()
521                .flat_map(|command| &command.args)
522                .any(|arg| arg == "--remote")
523        );
524        assert!(plans[3].commands[0].args.contains(&"-P".into()));
525    }
526
527    fn git(repository: &Path, args: &[&str]) -> String {
528        let output = Command::new("git")
529            .args(args)
530            .current_dir(repository)
531            .output()
532            .unwrap();
533        assert!(
534            output.status.success(),
535            "git {args:?}: {}",
536            String::from_utf8_lossy(&output.stderr)
537        );
538        String::from_utf8(output.stdout).unwrap().trim().into()
539    }
540
541    fn fixture(temp: &Path) -> (CheckpointExportSpec, PathBuf) {
542        let worker_root = temp.join("worker");
543        fs::create_dir_all(&worker_root).unwrap();
544        let harness_home = temp.join("codex");
545        let native = harness_home.join("sessions/2026/08/09");
546        fs::create_dir_all(&native).unwrap();
547        fs::write(native.join(format!("rollout-{NATIVE}.jsonl")), b"native").unwrap();
548        let workspace = temp.join("workspace");
549        let repository = workspace.join("app");
550        fs::create_dir_all(&repository).unwrap();
551        git(&repository, &["init"]);
552        git(&repository, &["config", "user.email", "hel@example.test"]);
553        git(&repository, &["config", "user.name", "Hel Test"]);
554        fs::write(repository.join("README.md"), b"hello").unwrap();
555        git(&repository, &["add", "."]);
556        git(&repository, &["commit", "-m", "base"]);
557        git(
558            &repository,
559            &[
560                "remote",
561                "add",
562                "origin",
563                "https://github.com/example/app.git",
564            ],
565        );
566        let base = git(&repository, &["rev-parse", "HEAD"]);
567        let output = worker_root.join("source.hel.zip");
568        (
569            CheckpointExportSpec {
570                protocol_version: CHECKPOINT_EXPORT_PROTOCOL_VERSION,
571                session: SessionManifest {
572                    id: SESSION.into(),
573                    title: "test".into(),
574                    harness_kind: HarnessKind::Codex,
575                    profile_id: "codex-1".into(),
576                    native_session_id: NATIVE.into(),
577                    created_at: "2026-08-09T00:00:00Z".into(),
578                    checkpointed_at: "2026-08-09T00:01:00Z".into(),
579                    hel_version: "0.1.0".into(),
580                    relay_version: "0.1.0".into(),
581                    adapter_version: "test".into(),
582                },
583                target: TargetManifest {
584                    template_id: "local".into(),
585                    target_kind: "podman".into(),
586                    details: Default::default(),
587                },
588                bundle: BundleManifest {
589                    id: "bundle".into(),
590                    primary_repository: "app".into(),
591                },
592                relay_root: worker_root,
593                harness_home,
594                workspace_root: workspace,
595                repositories: vec![CheckpointRepositorySpec {
596                    id: "app".into(),
597                    relative_destination: "app".into(),
598                    capture: CheckpointRepositoryCapture::DeltaFrom { base_commit: base },
599                    origin_override: None,
600                }],
601                canonical_session: CanonicalSessionSnapshot {
602                    event_frontier: 1,
603                    event_frontier_digest: "a".repeat(64),
604                    session: CanonicalSessionState {
605                        execution: CanonicalExecutionState::Idle,
606                        last_activity_at_ms: Some(1),
607                        session_title: Some("test".into()),
608                        configuration: Default::default(),
609                    },
610                    transcript: vec![CanonicalTranscriptItem {
611                        stable_id: "user-1".into(),
612                        position: 1,
613                        latest_content_event_ordinal: None,
614                        created_at_ms: 1,
615                        last_changed_at_ms: 1,
616                        body: CanonicalTranscriptBody::User {
617                            content: vec![json!({"type": "text", "text": "hello"})],
618                        },
619                    }],
620                    queued_prompts: vec![CanonicalQueuedPrompt {
621                        command_id: "queued-1".into(),
622                        kind: CanonicalQueuedCommandKind::Prompt,
623                        content: vec![json!({"type": "text", "text": "next"})],
624                        queued_at_ms: 2,
625                    }],
626                },
627                output_path: output.clone(),
628            },
629            output,
630        )
631    }
632
633    struct CopyExecutor {
634        source: PathBuf,
635        calls: RefCell<usize>,
636    }
637    impl CommandExecutor for CopyExecutor {
638        fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
639            *self.calls.borrow_mut() += 1;
640            fs::copy(
641                &self.source,
642                command.args.last().context("missing destination")?,
643            )?;
644            Ok(CommandOutput {
645                status: 0,
646                stdout: vec![],
647                stderr: vec![],
648            })
649        }
650    }
651
652    struct SshDockerTransferExecutor {
653        archive: Vec<u8>,
654        commands: RefCell<Vec<CommandSpec>>,
655        fail_download: bool,
656        fail_staging_cleanup: bool,
657    }
658
659    impl SshDockerTransferExecutor {
660        fn new(archive: Vec<u8>) -> Self {
661            Self {
662                archive,
663                commands: RefCell::new(Vec::new()),
664                fail_download: false,
665                fail_staging_cleanup: false,
666            }
667        }
668    }
669
670    impl CommandExecutor for SshDockerTransferExecutor {
671        fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
672            self.commands.borrow_mut().push(command.clone());
673            let remote = command.args.last().map(String::as_str).unwrap_or_default();
674            if command.program == "scp" {
675                if self.fail_download {
676                    return Ok(CommandOutput {
677                        status: 23,
678                        stdout: Vec::new(),
679                        stderr: b"scp unavailable".to_vec(),
680                    });
681                }
682                fs::write(
683                    command
684                        .args
685                        .last()
686                        .context("missing local checkpoint path")?,
687                    &self.archive,
688                )?;
689            }
690            if command.program == "ssh"
691                && remote.contains("'rm' '-f' '--' '.local/share/hel/transfers/")
692                && self.fail_staging_cleanup
693            {
694                return Ok(CommandOutput {
695                    status: 19,
696                    stdout: Vec::new(),
697                    stderr: b"staging cleanup unavailable".to_vec(),
698                });
699            }
700            Ok(CommandOutput {
701                status: 0,
702                stdout: Vec::new(),
703                stderr: Vec::new(),
704            })
705        }
706    }
707
708    fn ssh_docker_locator() -> TargetLocator {
709        TargetLocator::SshDocker {
710            ssh: ssh(),
711            container_id: mj_core::targets::resource_name(SESSION).unwrap(),
712        }
713    }
714
715    #[test]
716    fn export_and_transfer_only_gate_after_local_verification() {
717        let temp = tempfile::tempdir().unwrap();
718        let (spec, source) = fixture(temp.path());
719        let target = export_checkpoint(&spec).unwrap();
720        assert_eq!(target.event_frontier, 1);
721        assert_eq!(
722            target.event_frontier_digest,
723            spec.canonical_session.event_frontier_digest
724        );
725        let destination = temp.path().join("controller/session.hel.zip");
726        let locator = &locators()[0];
727        let gate = CheckpointTransfer {
728            locator,
729            session_id: SESSION,
730            operation_id: "test-transfer",
731            remote_archive: "/var/lib/hel/workers/source.hel.zip",
732            destination: &destination,
733            expected_sha256: &target.sha256,
734            expected_event_frontier: 1,
735            expected_event_frontier_digest: &spec.canonical_session.event_frontier_digest,
736        }
737        .execute(&CopyExecutor {
738            source,
739            calls: RefCell::new(0),
740        })
741        .unwrap();
742        assert!(gate.teardown_allowed());
743        assert_eq!(gate.event_frontier(), 1);
744        assert_eq!(
745            gate.event_frontier_digest(),
746            spec.canonical_session.event_frontier_digest
747        );
748        assert_eq!(
749            read_archive_verified(&destination).unwrap().archive_sha256,
750            gate.sha256()
751        );
752    }
753
754    /// The controller streams the export spec to save a round trip to the
755    /// target. Both spellings have to produce the same archive.
756    #[test]
757    fn a_streamed_spec_exports_the_same_archive_as_a_spec_file() {
758        let temp = tempfile::tempdir().unwrap();
759        let (mut spec, _) = fixture(temp.path());
760        let from_file = export_from_spec_file(&spec.output_path.with_extension("spec.json"))
761            .err()
762            .map(|error| format!("{error:#}"));
763        assert!(
764            from_file.is_some_and(|error| error.contains("read checkpoint export spec")),
765            "a missing spec file must still be reported as a read failure"
766        );
767
768        let spec_path = temp.path().join("checkpoint-spec.json");
769        spec.write(&spec_path).unwrap();
770        let from_file = export_from_spec_file(&spec_path).unwrap();
771        let file_archive = fs::read(&spec.output_path).unwrap();
772
773        spec.output_path = temp.path().join("worker/streamed.hel.zip");
774        let body = serde_json::to_vec(&spec).unwrap();
775        let streamed = export_from_spec_reader(&mut body.as_slice()).unwrap();
776        let streamed_archive = fs::read(&spec.output_path).unwrap();
777
778        assert_eq!(streamed.sha256, from_file.sha256);
779        assert_eq!(streamed.event_frontier, from_file.event_frontier);
780        assert_eq!(
781            streamed.event_frontier_digest,
782            from_file.event_frontier_digest
783        );
784        assert_eq!(streamed_archive, file_archive);
785        assert_eq!(
786            read_archive_verified(&spec.output_path)
787                .unwrap()
788                .archive_sha256,
789            streamed.sha256
790        );
791    }
792
793    #[test]
794    fn transfer_rejects_a_target_checksum_mismatch() {
795        let temp = tempfile::tempdir().unwrap();
796        let (spec, source) = fixture(temp.path());
797        export_checkpoint(&spec).unwrap();
798        let destination = temp.path().join("controller/session.hel.zip");
799        let unexpected_sha256 = "b".repeat(64);
800
801        let error = CheckpointTransfer {
802            locator: &locators()[0],
803            session_id: SESSION,
804            operation_id: "test-transfer",
805            remote_archive: "/var/lib/hel/workers/source.hel.zip",
806            destination: &destination,
807            expected_sha256: &unexpected_sha256,
808            expected_event_frontier: 1,
809            expected_event_frontier_digest: &spec.canonical_session.event_frontier_digest,
810        }
811        .execute(&CopyExecutor {
812            source,
813            calls: RefCell::new(0),
814        })
815        .unwrap_err();
816
817        assert!(format!("{error:#}").contains("checkpoint checksums differ"));
818        assert!(!destination.exists());
819    }
820
821    #[test]
822    fn ssh_docker_failed_hash_cleans_host_staging_but_preserves_container_archive() {
823        let temp = tempfile::tempdir().unwrap();
824        let destination = temp.path().join("controller/session.hel.zip");
825        let executor = SshDockerTransferExecutor::new(vec![b'x'; 128 * 1024]);
826        let error = CheckpointTransfer {
827            locator: &ssh_docker_locator(),
828            session_id: SESSION,
829            operation_id: "test-transfer",
830            remote_archive: "/var/lib/hel/workers/source.hel.zip",
831            destination: &destination,
832            expected_sha256: &"0".repeat(64),
833            expected_event_frontier: 1,
834            expected_event_frontier_digest: &"a".repeat(64),
835        }
836        .execute(&executor)
837        .unwrap_err();
838
839        assert!(format!("{error:#}").contains("checkpoint checksums differ"));
840        assert!(!destination.exists());
841        let commands = executor.commands.borrow();
842        assert!(commands.iter().any(|command| {
843            command.program == "ssh"
844                && command.args.last().is_some_and(|remote| {
845                    remote.contains("'rm' '-f' '--' '.local/share/hel/transfers/")
846                })
847        }));
848        assert!(!commands.iter().any(|command| {
849            command
850                .args
851                .last()
852                .is_some_and(|remote| remote.contains("'docker' 'exec'"))
853        }));
854    }
855
856    #[test]
857    fn ssh_docker_download_and_staging_cleanup_errors_keep_the_original_failure() {
858        let temp = tempfile::tempdir().unwrap();
859        let destination = temp.path().join("controller/session.hel.zip");
860        let mut executor = SshDockerTransferExecutor::new(vec![b'x'; 128 * 1024]);
861        executor.fail_download = true;
862        executor.fail_staging_cleanup = true;
863        let error = CheckpointTransfer {
864            locator: &ssh_docker_locator(),
865            session_id: SESSION,
866            operation_id: "test-transfer",
867            remote_archive: "/var/lib/hel/workers/source.hel.zip",
868            destination: &destination,
869            expected_sha256: &"0".repeat(64),
870            expected_event_frontier: 1,
871            expected_event_frontier_digest: &"a".repeat(64),
872        }
873        .execute(&executor)
874        .unwrap_err();
875
876        let text = format!("{error:#}");
877        assert!(text.contains("download target checkpoint"), "{text}");
878        assert!(text.contains("scp unavailable"), "{text}");
879        assert!(
880            text.contains("clean target checkpoint host staging also failed")
881                && text.contains("staging cleanup unavailable"),
882            "{text}"
883        );
884        let commands = executor.commands.borrow();
885        assert!(commands.iter().any(|command| {
886            command.program == "ssh"
887                && command.args.last().is_some_and(|remote| {
888                    remote.contains("'rm' '-f' '--' '.local/share/hel/transfers/")
889                })
890        }));
891        assert!(!commands.iter().any(|command| {
892            command
893                .args
894                .last()
895                .is_some_and(|remote| remote.contains("'docker' 'exec'"))
896        }));
897    }
898
899    #[test]
900    fn overlapping_remote_transfers_keep_their_own_bytes_and_cleanup() {
901        use std::collections::BTreeMap;
902        use std::sync::{Condvar, Mutex};
903        use std::time::Duration;
904
905        #[derive(Default)]
906        struct Staging {
907            files: BTreeMap<String, Vec<u8>>,
908            copies: usize,
909            first_cleaned: bool,
910        }
911        struct InterleavedExecutor<'a> {
912            staging: &'a (Mutex<Staging>, Condvar),
913            archive: &'a [u8],
914            first: bool,
915        }
916        impl CommandExecutor for InterleavedExecutor<'_> {
917            fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
918                let (mutex, changed) = self.staging;
919                let remote = command.args.last().unwrap();
920                // Fixture paths contain no quotes or spaces. Inspect the final
921                // remote argument, so a shared staging path really collides.
922                let remote_path = || {
923                    remote
924                        .rsplit(' ')
925                        .next()
926                        .unwrap()
927                        .trim_matches('\'')
928                        .to_owned()
929                };
930                match command.purpose.as_str() {
931                    "stage remote container checkpoint" => {
932                        let mut state = mutex.lock().unwrap();
933                        state.files.insert(remote_path(), self.archive.to_vec());
934                        state.copies += 1;
935                        changed.notify_all();
936                    }
937                    "download remote container checkpoint over SSH" => {
938                        let (state, timeout) = changed
939                            .wait_timeout_while(
940                                mutex.lock().unwrap(),
941                                Duration::from_secs(5),
942                                |state| state.copies < 2 || (!self.first && !state.first_cleaned),
943                            )
944                            .unwrap();
945                        ensure!(!timeout.timed_out(), "interleaved transfer stalled");
946                        let source = command.args[command.args.len() - 2]
947                            .split_once(':')
948                            .unwrap()
949                            .1;
950                        let bytes = state
951                            .files
952                            .get(source)
953                            .context("other transfer removed staging")?;
954                        fs::write(remote, bytes)?;
955                    }
956                    "remove remote checkpoint staging" => {
957                        let mut state = mutex.lock().unwrap();
958                        state.files.remove(&remote_path());
959                        if self.first {
960                            state.first_cleaned = true;
961                            changed.notify_all();
962                        }
963                    }
964                    "create remote checkpoint staging directory" => {}
965                    purpose => bail!("unexpected transfer command: {purpose}"),
966                }
967                Ok(CommandOutput {
968                    status: 0,
969                    stdout: Vec::new(),
970                    stderr: Vec::new(),
971                })
972            }
973        }
974
975        for locator in [locators().pop().unwrap(), ssh_docker_locator()] {
976            let directory = tempfile::tempdir().unwrap();
977            let staging = (Mutex::new(Staging::default()), Condvar::new());
978            let first = vec![b'a'; 192 * 1024];
979            let second = vec![b'b'; 256 * 1024];
980            std::thread::scope(|scope| {
981                let run = |operation, bytes: &[u8], is_first| {
982                    let source = directory.path().join(format!("{operation}-source.zip"));
983                    fs::write(&source, bytes).unwrap();
984                    let destination = directory.path().join(format!("{operation}-verified.zip"));
985                    let digest = checkpoint_sha256(&source).unwrap();
986                    let gate = CheckpointTransfer {
987                        locator: &locator,
988                        session_id: SESSION,
989                        operation_id: operation,
990                        remote_archive: &format!("/workers/{operation}.zip"),
991                        destination: &destination,
992                        expected_sha256: &digest,
993                        expected_event_frontier: 1,
994                        expected_event_frontier_digest: &"a".repeat(64),
995                    }
996                    .execute(&InterleavedExecutor {
997                        staging: &staging,
998                        archive: bytes,
999                        first: is_first,
1000                    })
1001                    .unwrap();
1002                    assert_eq!(fs::read(gate.archive_path()).unwrap(), bytes);
1003                    assert_eq!(gate.sha256(), digest);
1004                };
1005                let first_run = scope.spawn(move || run("first", &first, true));
1006                let second_run = scope.spawn(move || run("second", &second, false));
1007                first_run.join().unwrap();
1008                second_run.join().unwrap();
1009            });
1010            assert!(staging.0.lock().unwrap().files.is_empty());
1011        }
1012    }
1013
1014    #[test]
1015    fn corrupt_or_truncated_transfer_preserves_previous_checkpoint_and_reports_evidence() {
1016        let temp = tempfile::tempdir().unwrap();
1017        let source = temp.path().join("source.zip");
1018        let original = vec![b'a'; 192 * 1024];
1019        fs::write(&source, &original).unwrap();
1020        let expected_sha256 = checkpoint_sha256(&source).unwrap();
1021        let mut corrupt = original.clone();
1022        corrupt[80 * 1024] = b'b';
1023        for bytes in [corrupt, original[..1024].to_vec()] {
1024            let downloaded = temp.path().join("download.zip");
1025            fs::write(&downloaded, &bytes).unwrap();
1026            let actual_sha256 = checkpoint_sha256(&downloaded).unwrap();
1027            let destination = temp.path().join("previous.zip");
1028            fs::write(&destination, b"previous verified checkpoint").unwrap();
1029            let error = CheckpointTransfer {
1030                locator: &ssh_docker_locator(),
1031                session_id: SESSION,
1032                operation_id: "failed-export",
1033                remote_archive: "/workers/failed-export.zip",
1034                destination: &destination,
1035                expected_sha256: &expected_sha256,
1036                expected_event_frontier: 1,
1037                expected_event_frontier_digest: &"a".repeat(64),
1038            }
1039            .execute(&SshDockerTransferExecutor::new(bytes.clone()))
1040            .unwrap_err();
1041            let detail = format!("{error:#}");
1042            for expected in [
1043                "complete checkpoint archive",
1044                SESSION,
1045                "failed-export",
1046                &expected_sha256,
1047                &actual_sha256,
1048                &format!("downloaded_bytes={}", bytes.len()),
1049                "target archive retained at /workers/failed-export.zip",
1050                "retry a fresh export",
1051            ] {
1052                assert!(
1053                    detail.contains(expected),
1054                    "missing {expected:?} from {detail}"
1055                );
1056            }
1057            assert_eq!(
1058                fs::read(&destination).unwrap(),
1059                b"previous verified checkpoint"
1060            );
1061            assert_eq!(fs::read(&source).unwrap(), original);
1062        }
1063    }
1064}