Skip to main content

mj_controller/controller/worker_binary/
binary_select.rs

1use super::*;
2use mj_core::hex::lower_hex;
3
4/// Linux appends " (deleted)" to `/proc/<pid>/exe` for a removed image. That
5/// marker belongs in a message but never in a decision, which `is_file` makes.
6pub(super) fn display_path(path: &Path) -> String {
7    let text = path.to_string_lossy();
8    text.strip_suffix(" (deleted)").unwrap_or(&text).to_owned()
9}
10
11/// The architecture a configured template names outright, if it names one. A
12/// container `platform` such as `linux/arm64` decides what the target runs
13/// whatever the controller's own machine is, and it is the only architecture a
14/// configured target template can state: the configured `AwsEc2` variant names
15/// a launch template, whose instance type is only discoverable through the AWS
16/// API.
17pub(super) fn template_architecture(
18    template: &mj_core::config::TargetTemplate,
19) -> Option<&'static str> {
20    use mj_core::config::TargetTemplate as Template;
21    let platform = match template {
22        Template::LocalPodman { container }
23        | Template::LocalDocker { container }
24        | Template::AppleContainer { container }
25        | Template::SshPodman { container, .. }
26        | Template::SshDocker { container, .. } => container.platform.as_deref()?,
27        Template::LocalBare | Template::SshBare { .. } | Template::AwsEc2 { .. } => return None,
28    };
29    // Platform strings appear as "linux/arm64", "arm64", or "linux/arm64/v8".
30    platform
31        .split('/')
32        .find_map(|part| targets::normalize_architecture(part.trim()).ok())
33}
34
35/// Architectures a resume must be able to serve, knowing only the configured
36/// template. Provisioning learns the real answer by running `uname -m` on the
37/// live target; a resume has no target yet, so this uses what is knowable
38/// without one: an architecture the template names, else the controller's own
39/// architecture for a target that runs on this machine, else either Linux
40/// architecture for a remote target.
41pub(super) fn preflight_architectures(
42    template: &mj_core::config::TargetTemplate,
43) -> Vec<&'static str> {
44    use mj_core::config::TargetTemplate as Template;
45    if let Some(arch) = template_architecture(template) {
46        return vec![arch];
47    }
48    match template {
49        Template::LocalBare
50        | Template::LocalPodman { .. }
51        | Template::LocalDocker { .. }
52        | Template::AppleContainer { .. } => vec![std::env::consts::ARCH],
53        Template::SshBare { .. }
54        | Template::SshPodman { .. }
55        | Template::SshDocker { .. }
56        | Template::AwsEc2 { .. } => {
57            vec!["x86_64", "aarch64"]
58        }
59    }
60}
61
62/// Resolve the worker before expensive provisioning or transcript compaction.
63/// Existing SSH bare hosts are probed; disposable targets use their template.
64///
65/// A resume compacts a cross-harness transcript before it provisions anything,
66/// which costs minutes and paid model requests. Resolving the worker binary is
67/// performed before target creation, so a resume that could never install a worker
68/// must fail before spending any of that. Remote sources are downloaded and
69/// verified here too, before provisioning can create a container.
70pub(in crate::controller) fn preflight_worker_binary(
71    template: &mj_core::config::TargetTemplate,
72    executor: &impl CommandExecutor,
73) -> Result<()> {
74    if let mj_core::config::TargetTemplate::SshBare { ssh, .. } = template {
75        let command = targets::ssh_command(&SshTarget::from(ssh), ["uname", "-sm"])
76            .purpose("detect target platform");
77        let platform = probe_platform(executor, command)?;
78        return materialize_worker_source(worker_binary_for_arch(
79            platform.architecture,
80            WorkerBinaryRequirement::for_os(platform.os),
81        )?)
82        .map(|_| ());
83    }
84    // Existing SSH bare hosts were resolved above. Containers and the
85    // disposable EC2 backend require portable Linux workers.
86    let requirement = if matches!(template, mj_core::config::TargetTemplate::LocalBare) {
87        WorkerBinaryRequirement::LocalHost
88    } else {
89        WorkerBinaryRequirement::PortableLinux
90    };
91    let mut failure = None;
92    for arch in preflight_architectures(template) {
93        match worker_binary_for_arch(arch, requirement).and_then(materialize_worker_source) {
94            Ok(_) => return Ok(()),
95            Err(error) => failure = Some(error),
96        }
97    }
98    match failure {
99        // The message is the one provisioning would have printed later, so the
100        // user reads the same fix, sooner.
101        Some(error) => Err(error).context("preflight the worker binary before provisioning"),
102        None => Ok(()),
103    }
104}
105
106/// Why no worker binary could serve a session on this target, or `None` when
107/// one can (or when that cannot be told without a host that does not answer).
108///
109/// This is the resolution `preflight_worker_binary` performs, without the
110/// download or the build check, so the daemon can refuse a launch up front
111/// instead of admitting it and failing the session seconds later. An
112/// existing SSH host is asked its platform; a host that does not answer, or
113/// answers with a platform Mjolnir does not support, is not a worker-source
114/// problem and is left to the launch to report.
115pub(crate) fn worker_source_problem(
116    template: &mj_core::config::TargetTemplate,
117    executor: &impl CommandExecutor,
118) -> Option<String> {
119    if let mj_core::config::TargetTemplate::SshBare { ssh, .. } = template {
120        let command = targets::ssh_command(&SshTarget::from(ssh), ["uname", "-sm"])
121            .purpose("detect target platform");
122        let platform = probe_platform(executor, command).ok()?;
123        return worker_binary_for_arch(
124            platform.architecture,
125            WorkerBinaryRequirement::for_os(platform.os),
126        )
127        .err()
128        .map(|error| format!("{error:#}"));
129    }
130    let requirement = if matches!(template, mj_core::config::TargetTemplate::LocalBare) {
131        WorkerBinaryRequirement::LocalHost
132    } else {
133        WorkerBinaryRequirement::PortableLinux
134    };
135    let mut failure = None;
136    for arch in preflight_architectures(template) {
137        match worker_binary_for_arch(arch, requirement) {
138            Ok(_) => return None,
139            Err(error) => failure = Some(format!("{error:#}")),
140        }
141    }
142    failure
143}
144
145/// The worker an existing SSH host needs, with the target triple it is
146/// named by, read from the host's own platform. `None` for a template that
147/// is not a bare SSH host and for a host that does not say what it runs; the
148/// SSH reachability check reports the latter.
149pub fn ssh_worker_binary_prerequisite(
150    template: &mj_core::config::TargetTemplate,
151    executor: &impl CommandExecutor,
152) -> Option<(String, Result<WorkerBinaryAvailability>)> {
153    let mj_core::config::TargetTemplate::SshBare { ssh, .. } = template else {
154        return None;
155    };
156    let command = targets::ssh_command(&SshTarget::from(ssh), ["uname", "-sm"])
157        .purpose("detect target platform");
158    let platform = probe_platform(executor, command).ok()?;
159    let requirement = WorkerBinaryRequirement::for_os(platform.os);
160    Some((
161        requirement.triple(platform.architecture),
162        worker_binary_for_arch(platform.architecture, requirement),
163    ))
164}
165
166pub(in crate::controller) fn worker_binary_for(
167    locator: &targets::TargetLocator,
168    executor: &impl CommandExecutor,
169) -> Result<PathBuf> {
170    let platform = probe_platform(executor, targets::platform_probe(locator))?;
171    let requirement = if matches!(locator, targets::TargetLocator::LocalBare { .. }) {
172        WorkerBinaryRequirement::LocalHost
173    } else {
174        WorkerBinaryRequirement::for_os(platform.os)
175    };
176    materialize_worker_source(worker_binary_for_arch(platform.architecture, requirement)?)
177}
178
179fn materialize_worker_source(source: WorkerBinaryAvailability) -> Result<PathBuf> {
180    let path = match source {
181        WorkerBinaryAvailability::Local { path, .. } => Ok(path),
182        WorkerBinaryAvailability::Remote {
183            url,
184            sha256,
185            triple,
186        } => download_worker(&url, &sha256, &triple),
187    }?;
188    verify_worker_build(&path)?;
189    Ok(path)
190}
191
192fn probe_platform(
193    executor: &impl CommandExecutor,
194    command: CommandSpec,
195) -> Result<targets::TargetPlatform> {
196    let output = execute_checked(executor, command)?;
197    targets::TargetPlatform::parse(std::str::from_utf8(&output.stdout)?)
198}
199
200pub(super) fn download_worker(url: &str, expected_sha256: &str, triple: &str) -> Result<PathBuf> {
201    validate_worker_sha256(expected_sha256)?;
202    let digest = expected_sha256.to_ascii_lowercase();
203    let directory = data_dir().join("workers").join("pinned");
204    let destination = directory.join(&digest).join("hel");
205    std::fs::create_dir_all(destination.parent().unwrap_or(&directory))?;
206    if destination.is_file() {
207        let bytes = std::fs::read(&destination).with_context(|| {
208            format!(
209                "read cached worker for {triple} from {}",
210                destination.display()
211            )
212        })?;
213        if lower_hex(Sha256::digest(&bytes)).eq_ignore_ascii_case(expected_sha256) {
214            verify_worker_build(&destination)?;
215            return Ok(destination);
216        }
217        bail!(
218            "content-addressed worker cache {} does not match {} checksum",
219            destination.display(),
220            expected_sha256
221        );
222    }
223    let bytes = on_dedicated_thread(|| {
224        Ok(reqwest::blocking::Client::builder()
225            .timeout(std::time::Duration::from_secs(120))
226            .build()?
227            .get(url)
228            .send()?
229            .error_for_status()?
230            .bytes()?)
231    })?;
232    let actual = lower_hex(Sha256::digest(&bytes));
233    if !actual.eq_ignore_ascii_case(expected_sha256) {
234        bail!("downloaded worker checksum mismatch: expected {expected_sha256}, got {actual}");
235    }
236    std::fs::create_dir_all(&directory)?;
237    let mut temporary = tempfile::NamedTempFile::new_in(&directory)?;
238    std::io::Write::write_all(&mut temporary, &bytes)?;
239    temporary.as_file_mut().sync_all()?;
240    #[cfg(unix)]
241    {
242        use std::os::unix::fs::PermissionsExt;
243        std::fs::set_permissions(temporary.path(), std::fs::Permissions::from_mode(0o700))?;
244    }
245    publish_cached_worker(temporary, &directory, &digest)
246}
247
248pub(super) fn validate_worker_sha256(expected_sha256: &str) -> Result<()> {
249    if expected_sha256.len() != 64 || !expected_sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
250    {
251        bail!("MJ_WORKER_SHA256 must be a 64-character hexadecimal digest");
252    }
253    Ok(())
254}