Skip to main content

microvms_app/control/
ensure.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Content-addressed build-or-reuse for an image: what `Sandbox::ensure_image` does (#221).
3//!
4//! A harness with many concurrent trials of one task needs one image for them all: the one
5//! already built, the one a sibling trial is building, or — when there is none — one it
6//! builds itself. This module is that composition, from build inputs to a usable image:
7//!
8//! 1. **Name** (IMAGE-6). `<prefix>-<hash12>`, the hash over the artifact's inputs (the
9//!    daemon, the Dockerfile, the build context — [`artifact_content_hash_with_context`])
10//!    plus the base image, the size class and every field the service sets only when it
11//!    creates the image (identity repair, `inherit_workdir`, the hook timeouts, the log
12//!    destination), because an image is created on one base at one size with those fixed.
13//!    Equal inputs name one image; any changed input names a fresh one, which is what keeps a
14//!    stale snapshot, or one built without a field the caller asked for, from being served
15//!    under a reused name.
16//! 2. **Artifact** (IMAGE-7). The Dockerfile, the daemon, and the context, zipped with fixed
17//!    dates and modes, so equal inputs are one object at one content-addressed key.
18//! 3. **ARN and upload** (IMAGE-8). The service takes image ARNs, not names, so the ARN is
19//!    built from the caller's account, looked up once per sandbox
20//!    ([`super::services::BuildServices::caller_account`]). The artifact goes to
21//!    `s3://<bucket>/<prefix>/<name>/artifact.zip`, and only when a build is needed.
22//! 4. **Decide** (IMAGE-9, IMAGE-10). After each describe, [`plan`] says what to do: reuse a
23//!    ready image, wait for one building, delete a failed one (or, under `force`, any one)
24//!    and wait for the name to be free, or build.
25//! 5. **Race** (IMAGE-11). Two trials that both find the name free both create; the service
26//!    accepts one and refuses the other. The loser describes again and waits for the winner's
27//!    build, returning it only once it is ready, and marks it reused.
28//!
29//! `crates/model/src/image.rs` is this module as a Stateright model: two concurrent callers against
30//! a platform that moves the image through its states, checked over every interleaving. The
31//! decision table there is [`plan`]'s, row for row.
32//!
33//! # Everything local happens before the first call
34//!
35//! [`prepare`] derives the name, the key, the artifact, and the create request, and runs the
36//! create call's preflight, with zero calls. A request this client refuses costs nothing: not
37//! the account lookup, not a describe, not the upload.
38//!
39//! # Timeouts
40//!
41//! The build wait is the one `build` uses — [`WaitOpts::default`], 45 minutes with a
42//! four-minute stall probe — or the caller's `wait_timeout`. The wait for a deleted name to
43//! free up is [`DEFAULT_DELETE_TIMEOUT`].
44
45use std::collections::BTreeMap;
46use std::time::Duration;
47
48use super::artifact::{BaseImage, artifact_content_hash_with_context, build_artifact_with_context};
49use super::context::BuildContext;
50use super::image::{DEFAULT_BUILD_TIMEOUT, Image, WaitOpts};
51use super::services::BuildServices;
52use super::transport::{Call, paths, send_accepting};
53use super::{ControlPlane, CreateImageRequest, ops};
54use crate::error::{Error, ErrorKind};
55use crate::sizing::SizeClass;
56
57/// How long a deleted image may take to free its name before `ensure_image` gives up.
58///
59/// Five minutes, the bound a harness provider measured for the same wait; a deletion that
60/// has not finished by then is reported rather than waited on for the build's 45.
61pub const DEFAULT_DELETE_TIMEOUT: Duration = Duration::from_secs(300);
62
63/// The longest name the service's `ImageName` admits.
64const MAX_IMAGE_NAME: usize = 64;
65
66/// The hex characters of the identity hash in a name.
67const NAME_HASH_LEN: usize = 12;
68
69/// What a describe of the ensured image found.
70#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
71pub enum Found {
72    /// No image under the name.
73    Absent,
74    /// `CREATING` or `UPDATING`, or a state this client does not know, which is waited on.
75    Building,
76    /// `CREATED` or `UPDATED` (or a tolerated ready spelling): usable.
77    Ready,
78    /// `CREATE_FAILED`, `UPDATE_FAILED`, or `DELETE_FAILED`.
79    Failed,
80    /// `DELETING`.
81    Deleting,
82}
83
84impl Found {
85    /// The reading of a `MicrovmImageState`.
86    pub fn from_state(state: &str) -> Self {
87        if Image::is_ready(state) {
88            Found::Ready
89        } else if Image::is_failed(state) {
90            Found::Failed
91        } else if state == "DELETING" {
92            Found::Deleting
93        } else if state == "DELETED" {
94            Found::Absent
95        } else {
96            Found::Building
97        }
98    }
99}
100
101/// What `ensure_image` does next with what a describe found.
102#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
103pub enum Plan {
104    /// Return the ready image; nothing is uploaded or created.
105    Reuse,
106    /// Wait for the running build to settle.
107    Wait,
108    /// Delete the image, wait until the name is free, then build.
109    Delete,
110    /// A deletion is under way: wait until the name is free, then build.
111    AwaitAbsent,
112    /// Upload the artifact and create the image.
113    Build,
114}
115
116/// The decision `ensure_image` makes after each describe (IMAGE-9, IMAGE-10).
117///
118/// A ready image is reused unless the caller forces a rebuild; a building one is waited on
119/// either way, because the service refuses to delete an image in `CREATING` — a forced
120/// caller deletes what the build settles to. A failed image is deleted whether or not the
121/// caller forced, because the name is content-addressed: the only way to a fresh build under
122/// it is to free it first. `crates/model/src/image.rs` checks this table against two concurrent
123/// callers, and `the_plan_table` holds the two equal.
124pub fn plan(found: Found, force: bool) -> Plan {
125    match (found, force) {
126        (Found::Absent, _) => Plan::Build,
127        (Found::Ready, false) => Plan::Reuse,
128        (Found::Ready, true) | (Found::Failed, _) => Plan::Delete,
129        (Found::Building, _) => Plan::Wait,
130        (Found::Deleting, _) => Plan::AwaitAbsent,
131    }
132}
133
134/// The image name for a prefix and an identity hash: `<prefix>-<hash12>` (IMAGE-6).
135///
136/// The prefix is reduced to what the service's `ImageName` pattern admits
137/// (`[a-zA-Z0-9-_]`), every other character becoming `-`, trimmed of `-` at both ends, and
138/// cut so the whole name fits the pattern's 64 characters. A prefix with nothing left is
139/// refused: the name would be the hash alone, and nothing in the account would say what it
140/// is for.
141pub fn ensured_image_name(prefix: &str, identity_hash: &str) -> Result<String, Error> {
142    let sanitized: String = prefix
143        .chars()
144        .map(|c| {
145            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
146                c
147            } else {
148                '-'
149            }
150        })
151        .collect();
152    let room = MAX_IMAGE_NAME - 1 - NAME_HASH_LEN;
153    let cut: String = sanitized.trim_matches('-').chars().take(room).collect();
154    let cut = cut.trim_end_matches('-');
155    if cut.is_empty() {
156        return Err(Error::invalid_arg(format!(
157            "the image name prefix {prefix:?} has no character the service's ImageName pattern \
158             admits ([a-zA-Z0-9-_]), so the name would be a bare hash. Pass a prefix that says \
159             what the image is for, such as the task's name."
160        )));
161    }
162    let hash: String = identity_hash.chars().take(NAME_HASH_LEN).collect();
163    Ok(format!("{cut}-{hash}"))
164}
165
166/// The identity an ensured image is named for (IMAGE-6): the artifact's content hash, the
167/// base image's name (what `baseImageArn` names), and the size class's baseline, for an image
168/// created with [`CreateImageRequest::new`]'s create-only fields.
169///
170/// Length-prefixed and tagged, so no two field boundaries can collide.
171pub fn image_identity_hash(artifact_hash: &str, base: &BaseImage, size: SizeClass) -> String {
172    pinned_identity_hash(artifact_hash, base, None, size)
173}
174
175/// [`image_identity_hash`] with the base version a build pins (#258): a pinned base builds
176/// on that version whatever the service's default later becomes, so the pin is identity.
177///
178/// The create-only fields are [`CreateImageRequest::new`]'s defaults; an ensure whose request
179/// sets others names its image for those (#280).
180pub fn pinned_identity_hash(
181    artifact_hash: &str,
182    base: &BaseImage,
183    base_version: Option<&str>,
184    size: SizeClass,
185) -> String {
186    let defaults = CreateImageRequest::new("", Vec::new(), "", "");
187    identity_hash(
188        artifact_hash,
189        base,
190        base_version,
191        size,
192        &CreateOnly::of(&defaults),
193    )
194}
195
196/// The identity of the image `create` would create, over the artifact `artifact_hash` names:
197/// what [`prepare`] names an ensured image for, and what the agent recipe's two-step path
198/// names the same image for.
199pub(crate) fn create_identity_hash(artifact_hash: &str, create: &CreateImageRequest) -> String {
200    identity_hash(
201        artifact_hash,
202        &create.base_image,
203        create.base_image_version.as_deref(),
204        create.size,
205        &CreateOnly::of(create),
206    )
207}
208
209/// The fields of a created image the service sets at `CreateMicrovmImage` and no later call
210/// changes (#280). They are identity because a reuse returns the image as it was created: an
211/// ensure that left them out of the name could hand a caller asking for identity repair an
212/// image built without it. Tags are not here, because a reused image keeps its own.
213struct CreateOnly<'a> {
214    repair_guest_identity: bool,
215    inherit_workdir: bool,
216    run_hook_timeout: u32,
217    build_hook_timeout: u32,
218    log_group: Option<&'a str>,
219    log_stream: Option<&'a str>,
220}
221
222impl<'a> CreateOnly<'a> {
223    fn of(create: &'a CreateImageRequest) -> Self {
224        Self {
225            repair_guest_identity: create.repair_guest_identity,
226            inherit_workdir: create.inherit_workdir,
227            run_hook_timeout: create.run_hook_timeout.as_secs(),
228            build_hook_timeout: create.build_hook_timeout.as_secs(),
229            log_group: create.log_group.as_deref(),
230            log_stream: create.log_stream.as_deref(),
231        }
232    }
233}
234
235/// The one identity stream. Every field is length-prefixed and each optional or create-only
236/// one is tagged, so no two field boundaries can collide and an absent value never reads as an
237/// empty one.
238///
239/// The domain tag is `microvms-ensure-image/2`: adding the create-only fields renames every
240/// ensured image once (#280), and the version says which stream a name came from. An image
241/// named by `/1` is never reused, which is the point: it may have been created without a field
242/// its next caller asks for.
243fn identity_hash(
244    artifact_hash: &str,
245    base: &BaseImage,
246    base_version: Option<&str>,
247    size: SizeClass,
248    create: &CreateOnly<'_>,
249) -> String {
250    use sha2::{Digest as _, Sha256};
251
252    fn put(hasher: &mut Sha256, field: &[u8]) {
253        hasher.update((field.len() as u64).to_be_bytes());
254        hasher.update(field);
255    }
256    fn put_optional(hasher: &mut Sha256, tag: &[u8], value: Option<&str>) {
257        put(hasher, tag);
258        match value {
259            None => put(hasher, &[0]),
260            Some(value) => {
261                put(hasher, &[1]);
262                put(hasher, value.as_bytes());
263            }
264        }
265    }
266
267    let mut hasher = Sha256::new();
268    put(&mut hasher, b"microvms-ensure-image/2");
269    put(&mut hasher, artifact_hash.as_bytes());
270    put(&mut hasher, base.name.as_bytes());
271    put(&mut hasher, &size.baseline_mib().to_be_bytes());
272    put_optional(&mut hasher, b"baseImageVersion", base_version);
273    for (tag, on) in [
274        (
275            b"repairGuestIdentity".as_slice(),
276            create.repair_guest_identity,
277        ),
278        (b"inheritWorkdir".as_slice(), create.inherit_workdir),
279    ] {
280        put(&mut hasher, tag);
281        put(&mut hasher, &[u8::from(on)]);
282    }
283    for (tag, secs) in [
284        (b"microvmHooks".as_slice(), create.run_hook_timeout),
285        (b"microvmImageHooks".as_slice(), create.build_hook_timeout),
286    ] {
287        put(&mut hasher, tag);
288        put(&mut hasher, &secs.to_be_bytes());
289    }
290    put_optional(&mut hasher, b"logGroup", create.log_group);
291    put_optional(&mut hasher, b"logStream", create.log_stream);
292    const_hex::encode(hasher.finalize())
293}
294
295/// The artifact's S3 key: `<prefix>/<name>/artifact.zip`, with the prefix's empty segments
296/// dropped so no key starts with or doubles a `/` (IMAGE-8).
297pub fn artifact_key(key_prefix: Option<&str>, name: &str) -> String {
298    let prefix: Vec<&str> = key_prefix
299        .unwrap_or("")
300        .split('/')
301        .filter(|segment| !segment.is_empty())
302        .collect();
303    if prefix.is_empty() {
304        format!("{name}/artifact.zip")
305    } else {
306        format!("{}/{name}/artifact.zip", prefix.join("/"))
307    }
308}
309
310/// The bucket and key of an `s3://<bucket>/<key>` URI, refused when it isn't one S3 admits.
311///
312/// For the upload of a build whose artifact URI the caller chose ([`crate::sandbox::Sandbox::
313/// upload_artifact_for`]); an ensure derives its own.
314pub fn s3_location(uri: &str) -> Result<(&str, &str), Error> {
315    let (bucket, key) = uri
316        .strip_prefix("s3://")
317        .and_then(|rest| rest.split_once('/'))
318        .filter(|(_, key)| !key.is_empty())
319        .ok_or_else(|| {
320            Error::invalid_arg(format!(
321                "the artifact URI {uri:?} is not s3://<bucket>/<key>, the only place \
322                 CreateMicrovmImage reads an artifact from."
323            ))
324        })?;
325    require_valid_bucket(bucket)?;
326    Ok((bucket, key))
327}
328
329/// Refuses a bucket name S3 does not admit, before any call.
330fn require_valid_bucket(bucket: &str) -> Result<(), Error> {
331    let chars_ok = bucket
332        .bytes()
333        .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'.' || b == b'-');
334    let ends_ok = bucket
335        .bytes()
336        .next()
337        .zip(bucket.bytes().last())
338        .is_some_and(|(first, last)| first.is_ascii_alphanumeric() && last.is_ascii_alphanumeric());
339    if (3..=63).contains(&bucket.len()) && chars_ok && ends_ok && !bucket.contains("..") {
340        return Ok(());
341    }
342    Err(Error::invalid_arg(format!(
343        "the S3 bucket {bucket:?} is not a bucket name S3 admits: 3 to 63 characters of \
344         lowercase letters, digits, `.` and `-`, starting and ending with a letter or digit. \
345         Pass the bucket the build role can read, in the sandbox's region."
346    )))
347}
348
349/// Refuses a key prefix that could not form a legal S3 key.
350fn require_valid_key_prefix(prefix: Option<&str>, name: &str) -> Result<(), Error> {
351    let Some(prefix) = prefix else {
352        return Ok(());
353    };
354    if prefix.chars().any(char::is_control) {
355        return Err(Error::invalid_arg(format!(
356            "the S3 key prefix {prefix:?} carries a control character, which no S3 key \
357             should; pass a plain path such as `harbor/images`."
358        )));
359    }
360    let key = artifact_key(Some(prefix), name);
361    if key.len() > 1024 {
362        return Err(Error::invalid_arg(format!(
363            "the S3 key prefix makes a {}-byte key, over S3's 1024; shorten it.",
364            key.len()
365        )));
366    }
367    Ok(())
368}
369
370/// Everything `Sandbox::ensure_image` needs.
371#[derive(Clone, Debug)]
372pub struct EnsureImageRequest {
373    /// The stem of the image name; the name is `<prefix>-<hash12>`.
374    pub name_prefix: String,
375    /// The daemon binary's bytes.
376    pub binary: Vec<u8>,
377    /// The Dockerfile, usually `wrap_dockerfile`'s output. Checked by the create call's
378    /// preflight before anything else happens.
379    pub dockerfile: String,
380    /// The files the Dockerfile may `COPY` ([`super::context::from_dir`]), or `None`.
381    pub context: Option<BuildContext>,
382    /// The bucket the artifact is uploaded to, in the sandbox's region.
383    pub s3_bucket: String,
384    /// A key prefix inside the bucket, or `None` for the bucket root.
385    pub s3_key_prefix: Option<String>,
386    /// The build role, which must read the bucket and write logs under
387    /// `/aws/lambda-microvms/`.
388    pub build_role_arn: String,
389    /// The size class the image is created with; part of the name.
390    pub size: SizeClass,
391    /// The base image, or `None` for [`BaseImage::from_dockerfile`] of the Dockerfile.
392    pub base_image: Option<BaseImage>,
393    /// Delete what exists under the name and build afresh.
394    pub force: bool,
395    /// Tags for a created image. Not part of the name: a reused image keeps its own, and
396    /// tags are not fixed at create.
397    pub tags: BTreeMap<String, String>,
398    /// The build wait's deadline, or `None` for [`DEFAULT_BUILD_TIMEOUT`].
399    pub wait_timeout: Option<Duration>,
400    /// A project's manifest and lockfile, baked as an environment layer (#74). Part of the
401    /// name: they enter the artifact.
402    pub project_files: Option<super::artifact::ProjectFiles>,
403    /// `baseImageVersion`, or `None` for the service's default. Part of the name when set.
404    pub base_image_version: Option<String>,
405    /// `logging.cloudWatch.logGroup` for a created image. Part of the name (#280): the
406    /// service sets an image's logging when it creates it, so a reused image logs where its
407    /// creator asked.
408    pub log_group: Option<String>,
409    /// The log-stream prefix for a created image; see [`CreateImageRequest::log_stream`].
410    /// Part of the name, as `log_group` is.
411    pub log_stream: Option<String>,
412    /// [`CreateImageRequest::repair_guest_identity`] for a created image. Part of the name
413    /// (#280): a reused image built without the repair would answer a caller who asked for it.
414    pub repair_guest_identity: bool,
415    /// [`CreateImageRequest::inherit_workdir`] for a created image. Part of the name.
416    pub inherit_workdir: bool,
417    /// The run-family hook timeouts for a created image. Part of the name.
418    pub run_hook_timeout: crate::hooks::RunHookTimeout,
419    /// The build-family hook timeouts for a created image. Part of the name.
420    pub build_hook_timeout: crate::hooks::BuildHookTimeout,
421}
422
423impl EnsureImageRequest {
424    /// A request with the defaults: no context, no key prefix, the default size class, the
425    /// base derived from the Dockerfile, no force, no tags, the default wait, and the create
426    /// request's own defaults for everything a created image carries.
427    pub fn new(
428        name_prefix: impl Into<String>,
429        binary: Vec<u8>,
430        dockerfile: impl Into<String>,
431        s3_bucket: impl Into<String>,
432        build_role_arn: impl Into<String>,
433    ) -> Self {
434        // The create request's defaults, read off it rather than restated, so a created
435        // image carries what `build_image` would have given it.
436        let defaults = CreateImageRequest::new("", Vec::new(), "", "");
437        Self {
438            name_prefix: name_prefix.into(),
439            binary,
440            dockerfile: dockerfile.into(),
441            context: None,
442            s3_bucket: s3_bucket.into(),
443            s3_key_prefix: None,
444            build_role_arn: build_role_arn.into(),
445            size: SizeClass::DEFAULT,
446            base_image: None,
447            force: false,
448            tags: BTreeMap::new(),
449            wait_timeout: None,
450            project_files: None,
451            base_image_version: None,
452            log_group: None,
453            log_stream: None,
454            repair_guest_identity: defaults.repair_guest_identity,
455            inherit_workdir: defaults.inherit_workdir,
456            run_hook_timeout: defaults.run_hook_timeout,
457            build_hook_timeout: defaults.build_hook_timeout,
458        }
459    }
460
461    /// The ensure request for a create request's inputs (#258): its name becomes the prefix
462    /// and every field it carries a created image carries too. `dockerfile` is the text the
463    /// build takes, which a caller with no Dockerfile of its own gets from
464    /// [`ControlPlane::dockerfile_for`]. The artifact URI and token label are dropped: the
465    /// ensure derives both from the name it computes.
466    pub fn from_create(
467        create: CreateImageRequest,
468        dockerfile: String,
469        s3_bucket: impl Into<String>,
470    ) -> Self {
471        let mut request = Self::new(
472            create.name,
473            create.binary,
474            dockerfile,
475            s3_bucket,
476            create.build_role_arn,
477        );
478        request.size = create.size;
479        request.base_image = Some(create.base_image);
480        request.tags = create.tags;
481        request.project_files = create.project_files;
482        request.base_image_version = create.base_image_version;
483        request.log_group = create.log_group;
484        request.log_stream = create.log_stream;
485        request.repair_guest_identity = create.repair_guest_identity;
486        request.inherit_workdir = create.inherit_workdir;
487        request.run_hook_timeout = create.run_hook_timeout;
488        request.build_hook_timeout = create.build_hook_timeout;
489        request
490    }
491}
492
493/// What `Sandbox::ensure_image` returns.
494#[derive(Clone, Debug)]
495pub struct EnsuredImage {
496    /// The ready image.
497    pub image: Image,
498    /// True when this call's own create did not build the image: it was ready, or a build
499    /// already running was waited out, or a sibling won the create race.
500    pub reused: bool,
501    /// Where the artifact is, whether or not this call uploaded it.
502    pub artifact_uri: String,
503    /// Whether this call uploaded the artifact.
504    pub uploaded: bool,
505    /// What reading the build context skipped, one line each ([`BuildContext::warnings`]).
506    pub warnings: Vec<String>,
507}
508
509/// Everything [`prepare`] decided locally: the name, the key, the artifact, and the create
510/// request, with the create's preflight already passed.
511#[derive(Clone, Debug)]
512pub struct Prepared {
513    /// `<prefix>-<hash12>`.
514    pub name: String,
515    pub bucket: String,
516    pub key: String,
517    /// The zip to upload when a build is needed.
518    pub artifact: Vec<u8>,
519    /// The create request, named and pointed at the artifact.
520    pub create: CreateImageRequest,
521    pub force: bool,
522    pub wait_timeout: Option<Duration>,
523    pub warnings: Vec<String>,
524}
525
526impl Prepared {
527    /// `s3://<bucket>/<key>`.
528    pub fn artifact_uri(&self) -> String {
529        format!("s3://{}/{}", self.bucket, self.key)
530    }
531}
532
533/// The local half of `ensure_image`: every name, key, and artifact decision and every guard,
534/// with zero calls.
535pub fn prepare(control: &ControlPlane, request: EnsureImageRequest) -> Result<Prepared, Error> {
536    require_valid_bucket(&request.s3_bucket)?;
537    let base = match request.base_image {
538        Some(base) => base,
539        None => BaseImage::from_dockerfile(&request.dockerfile)?,
540    };
541    let artifact_hash = artifact_content_hash_with_context(
542        &request.binary,
543        &request.dockerfile,
544        request.project_files.as_ref(),
545        request.context.as_ref(),
546    );
547    // The create request first, unnamed, so the name is the identity of exactly the image it
548    // creates (#280): every field it carries is read from it rather than restated.
549    let mut create = CreateImageRequest::new(
550        String::new(),
551        request.binary,
552        String::new(),
553        request.build_role_arn,
554    );
555    create.base_image = base;
556    create.dockerfile = Some(request.dockerfile);
557    create.size = request.size;
558    create.tags = request.tags;
559    create.project_files = request.project_files;
560    create.base_image_version = request.base_image_version;
561    create.log_group = request.log_group;
562    create.log_stream = request.log_stream;
563    create.repair_guest_identity = request.repair_guest_identity;
564    create.inherit_workdir = request.inherit_workdir;
565    create.run_hook_timeout = request.run_hook_timeout;
566    create.build_hook_timeout = request.build_hook_timeout;
567
568    let name = ensured_image_name(
569        &request.name_prefix,
570        &create_identity_hash(&artifact_hash, &create),
571    )?;
572    require_valid_key_prefix(request.s3_key_prefix.as_deref(), &name)?;
573    let key = artifact_key(request.s3_key_prefix.as_deref(), &name);
574    create.code_artifact_uri = format!("s3://{}/{key}", request.s3_bucket);
575    create.token_scope = Some(name.clone());
576    create.name = name.clone();
577    control.preflight(&create)?;
578
579    let artifact = build_artifact_with_context(
580        &create.binary,
581        create.dockerfile.as_deref().unwrap_or_default(),
582        create.project_files.as_ref(),
583        request.context.as_ref(),
584    )?;
585    Ok(Prepared {
586        name,
587        bucket: request.s3_bucket,
588        key,
589        artifact,
590        create,
591        force: request.force,
592        wait_timeout: request.wait_timeout,
593        warnings: request
594            .context
595            .map(|context| context.warnings().to_vec())
596            .unwrap_or_default(),
597    })
598}
599
600impl ControlPlane {
601    /// The image under `identifier`, or `None` when the service answers that there is none.
602    ///
603    /// `GetMicrovmImage`'s 404 is an answer here rather than an error: "no image under this
604    /// name" is the state `ensure_image` builds from. Every other failure is an error as
605    /// usual, and throttles are retried.
606    pub async fn describe_image(
607        &self,
608        identifier: &str,
609    ) -> Result<Option<ops::GetMicrovmImageResponseWire>, Error> {
610        super::require_valid_identifier("imageIdentifier", identifier)?;
611        let call = Call::get("GetMicrovmImage", paths::microvm_image(identifier));
612        let reply = send_accepting(self.transport(), call, &[404]).await?;
613        if reply.status == 404 {
614            return Ok(None);
615        }
616        Ok(Some(reply.json("GetMicrovmImage")?))
617    }
618
619    /// Polls until `identifier` is no longer building, answering what it settled to.
620    async fn await_settled(&self, identifier: &str, opts: WaitOpts) -> Result<Found, Error> {
621        let started = self.clock().elapsed();
622        loop {
623            let found = self
624                .describe_image(identifier)
625                .await?
626                .map_or(Found::Absent, |image| Found::from_state(&image.state));
627            if found != Found::Building {
628                return Ok(found);
629            }
630            let elapsed = self.clock().elapsed().saturating_sub(started);
631            if elapsed >= opts.timeout {
632                return Err(super::timed_out(
633                    &format!("image {identifier} did not finish building"),
634                    elapsed,
635                ));
636            }
637            self.clock().sleep(opts.poll_interval).await;
638        }
639    }
640
641    /// Deletes `identifier` and polls until its name is free (IMAGE-10).
642    ///
643    /// The delete is [`ControlPlane::delete_image`]'s, versions first. The poll ends when the
644    /// describe answers 404, or when the image under the name is no longer the one deleted —
645    /// a sibling rebuilt it, and the create that follows will join that build — or with
646    /// [`ErrorKind::Platform`] for `DELETE_FAILED` and a timeout at
647    /// [`DEFAULT_DELETE_TIMEOUT`].
648    async fn delete_and_await_absent(
649        &self,
650        identifier: &str,
651        deleted_state: &str,
652        poll_interval: Duration,
653    ) -> Result<(), Error> {
654        self.delete_image(identifier, 3, poll_interval).await;
655        self.await_absent(identifier, Some(deleted_state), poll_interval)
656            .await
657    }
658
659    /// Polls until the name is free, treating `deleted_state` — the state a delete was just
660    /// issued against — as a deletion the describe has not caught up with yet.
661    async fn await_absent(
662        &self,
663        identifier: &str,
664        deleted_state: Option<&str>,
665        poll_interval: Duration,
666    ) -> Result<(), Error> {
667        let started = self.clock().elapsed();
668        loop {
669            let Some(image) = self.describe_image(identifier).await? else {
670                return Ok(());
671            };
672            if image.state == "DELETE_FAILED" {
673                return Err(Error::new(
674                    ErrorKind::Platform,
675                    format!(
676                        "image {identifier} reports DELETE_FAILED, so its name cannot be \
677                         freed for a rebuild. Delete it by hand (`aws lambda-microvms \
678                         delete-microvm-image`), or pass a different name prefix."
679                    ),
680                ));
681            }
682            let pending = image.state == "DELETING" || Some(image.state.as_str()) == deleted_state;
683            if !pending {
684                return Ok(());
685            }
686            let elapsed = self.clock().elapsed().saturating_sub(started);
687            if elapsed >= DEFAULT_DELETE_TIMEOUT {
688                return Err(super::timed_out(
689                    &format!("image {identifier} was still present after its deletion"),
690                    elapsed,
691                ));
692            }
693            self.clock().sleep(poll_interval).await;
694        }
695    }
696
697    /// Whether the name is free or being freed: the image a wait was on has gone.
698    async fn vanished(&self, identifier: &str) -> Result<bool, Error> {
699        Ok(match self.describe_image(identifier).await? {
700            None => true,
701            Some(image) => Found::from_state(&image.state) == Found::Deleting,
702        })
703    }
704}
705
706fn image_from_wire(image: ops::GetMicrovmImageResponseWire, size: SizeClass) -> Image {
707    Image {
708        identifier: image.image_arn,
709        name: image.name,
710        version: image.latest_active_image_version.unwrap_or_default(),
711        state: image.state,
712        size,
713        log_group: None,
714        log_stream: None,
715    }
716}
717
718/// Where one `ensure_image` call is, across its describes.
719#[derive(Default)]
720struct Progress {
721    uploaded: bool,
722    /// This call deleted what it found once; a forced rebuild is one-shot.
723    deleted: bool,
724    /// This call went back to the describe after the image it waited on disappeared.
725    redescribed: bool,
726}
727
728/// The remote half of `ensure_image`: describe, decide, and build, reuse, wait, or join.
729pub(crate) async fn ensure(
730    control: &ControlPlane,
731    services: &dyn BuildServices,
732    account: &str,
733    prepared: Prepared,
734) -> Result<EnsuredImage, Error> {
735    let arn = format!(
736        "arn:aws:lambda:{}:{account}:microvm-image:{}",
737        control.region().as_str(),
738        prepared.name
739    );
740    let size = prepared.create.size;
741    let wait = WaitOpts {
742        timeout: prepared.wait_timeout.unwrap_or(DEFAULT_BUILD_TIMEOUT),
743        ..WaitOpts::default()
744    };
745    let done = |image: Image, reused: bool, uploaded: bool| EnsuredImage {
746        image,
747        reused,
748        artifact_uri: prepared.artifact_uri(),
749        uploaded,
750        warnings: prepared.warnings.clone(),
751    };
752    let mut progress = Progress::default();
753
754    loop {
755        let described = control.describe_image(&arn).await?;
756        let found = described
757            .as_ref()
758            .map_or(Found::Absent, |image| Found::from_state(&image.state));
759        let force = prepared.force && !progress.deleted;
760        match plan(found, force) {
761            Plan::Reuse => {
762                let image = described.expect("a ready image was described");
763                return Ok(done(image_from_wire(image, size), true, progress.uploaded));
764            }
765            Plan::Wait if force => {
766                // The service refuses to delete a build in progress: let it settle, then
767                // delete whatever it settled to.
768                let settled = control.await_settled(&arn, wait).await?;
769                if settled != Found::Absent {
770                    let state = control
771                        .describe_image(&arn)
772                        .await?
773                        .map(|image| image.state)
774                        .unwrap_or_default();
775                    control
776                        .delete_and_await_absent(&arn, &state, wait.poll_interval)
777                        .await?;
778                }
779                progress.deleted = true;
780            }
781            Plan::Wait => match control.wait_for_image(&arn, size, wait).await {
782                Ok(image) => return Ok(done(image, true, progress.uploaded)),
783                Err(error) => {
784                    if !progress.redescribed && control.vanished(&arn).await? {
785                        progress.redescribed = true;
786                        continue;
787                    }
788                    return Err(error);
789                }
790            },
791            Plan::Delete => {
792                let state = described.map(|image| image.state).unwrap_or_default();
793                control
794                    .delete_and_await_absent(&arn, &state, wait.poll_interval)
795                    .await?;
796                progress.deleted = true;
797            }
798            Plan::AwaitAbsent => control.await_absent(&arn, None, wait.poll_interval).await?,
799            Plan::Build => {}
800        }
801
802        // The name is free as far as this call knows: upload, create, and wait.
803        if !progress.uploaded {
804            services
805                .put_object(&prepared.bucket, &prepared.key, prepared.artifact.clone())
806                .await?;
807            progress.uploaded = true;
808        }
809        match control.create_image(prepared.create.clone()).await {
810            Ok(created) => match control
811                .wait_for_image(&created.identifier, size, wait)
812                .await
813            {
814                Ok(mut built) => {
815                    // The logging config survives the wait from the create's own record, as
816                    // in `Sandbox::build_image`: the readback cannot carry it.
817                    built.log_group = created.log_group;
818                    built.log_stream = created.log_stream;
819                    return Ok(done(built, false, true));
820                }
821                Err(error) => {
822                    if !progress.redescribed && control.vanished(&arn).await? {
823                        progress.redescribed = true;
824                        continue;
825                    }
826                    return Err(error);
827                }
828            },
829            // IMAGE-11: a refused create is a race lost when something now holds the name.
830            Err(refused) if refused.kind() == ErrorKind::Platform => {
831                match control.describe_image(&arn).await? {
832                    Some(image) if Found::from_state(&image.state) == Found::Ready => {
833                        return Ok(done(image_from_wire(image, size), true, true));
834                    }
835                    Some(image)
836                        if matches!(
837                            Found::from_state(&image.state),
838                            Found::Building | Found::Failed
839                        ) =>
840                    {
841                        // Waits for the winner's build; a failed build answers its own
842                        // diagnosis rather than this call's refused create.
843                        match control.wait_for_image(&arn, size, wait).await {
844                            Ok(image) => return Ok(done(image, true, true)),
845                            Err(error) => {
846                                if !progress.redescribed && control.vanished(&arn).await? {
847                                    progress.redescribed = true;
848                                    continue;
849                                }
850                                return Err(error);
851                            }
852                        }
853                    }
854                    _ if !progress.redescribed => {
855                        progress.redescribed = true;
856                        continue;
857                    }
858                    _ => return Err(refused),
859                }
860            }
861            Err(error) => return Err(error),
862        }
863    }
864}
865
866#[cfg(test)]
867mod tests {
868    use std::sync::atomic::{AtomicUsize, Ordering};
869    use std::sync::{Arc, Mutex};
870
871    use super::*;
872    use crate::control::fake::{self, Answer, FakeControlPlane, TestClock};
873    use crate::region::Region;
874    use crate::sandbox::Sandbox;
875    use crate::testing::aarch64_daemon;
876
877    const ACCOUNT: &str = "123456789012";
878    const BUCKET: &str = "artifact-bucket";
879    const ROLE: &str = "arn:aws:iam::123456789012:role/build";
880
881    /// STS and S3, recorded.
882    #[derive(Default)]
883    struct FakeServices {
884        account_calls: AtomicUsize,
885        puts: Mutex<Vec<(String, String, Vec<u8>)>>,
886    }
887
888    impl BuildServices for FakeServices {
889        fn caller_account(&self) -> futures_util::future::BoxFuture<'_, Result<String, Error>> {
890            self.account_calls.fetch_add(1, Ordering::SeqCst);
891            Box::pin(async { Ok(ACCOUNT.to_string()) })
892        }
893
894        fn put_object<'a>(
895            &'a self,
896            bucket: &'a str,
897            key: &'a str,
898            bytes: Vec<u8>,
899        ) -> futures_util::future::BoxFuture<'a, Result<(), Error>> {
900            self.puts.lock().expect("not poisoned").push((
901                bucket.to_string(),
902                key.to_string(),
903                bytes,
904            ));
905            Box::pin(async { Ok(()) })
906        }
907    }
908
909    fn sandbox(fake: &Arc<FakeControlPlane>, services: &Arc<FakeServices>) -> Sandbox {
910        let plane = crate::testing::control_plane(
911            fake.clone(),
912            Region::UsEast1,
913            Arc::new(TestClock::new()),
914        );
915        Sandbox::with_control_plane(plane).with_build_services(services.clone())
916    }
917
918    fn dockerfile() -> String {
919        super::super::artifact::wrap_dockerfile(
920            "FROM python:3.12-slim\nWORKDIR /app\n",
921            &super::super::artifact::WrapOptions::default(),
922        )
923        .expect("wraps")
924    }
925
926    fn request() -> EnsureImageRequest {
927        let mut request = EnsureImageRequest::new(
928            "task",
929            aarch64_daemon(b"daemon"),
930            dockerfile(),
931            BUCKET,
932            ROLE,
933        );
934        request.s3_key_prefix = Some("harbor/images".to_string());
935        request
936    }
937
938    /// The name `request()` derives, computed from the public pieces.
939    fn expected_name() -> String {
940        let request = request();
941        let base = BaseImage::from_dockerfile(&request.dockerfile).expect("a FROM");
942        let hash = super::super::artifact::artifact_content_hash_with_context(
943            &request.binary,
944            &request.dockerfile,
945            None,
946            None,
947        );
948        ensured_image_name("task", &image_identity_hash(&hash, &base, request.size))
949            .expect("a legal name")
950    }
951
952    /// `GetMicrovmImageOutput` for a usable image.
953    fn ready(name: &str) -> String {
954        format!(
955            r#"{{"imageArn": "arn:aws:lambda:us-east-1:{ACCOUNT}:microvm-image:{name}",
956                "name": "{name}", "state": "CREATED", "latestActiveImageVersion": "1",
957                "createdAt": 1754524800, "updatedAt": 1754528400, "tags": {{}}}}"#
958        )
959    }
960
961    fn state(name: &str, state: &str) -> Answer {
962        Answer::ok(fake::get_image_response(name, state))
963    }
964
965    fn absent() -> Answer {
966        Answer::failure(404, "Image not found")
967    }
968
969    fn uploads(services: &FakeServices) -> Vec<(String, String)> {
970        services
971            .puts
972            .lock()
973            .expect("not poisoned")
974            .iter()
975            .map(|(bucket, key, _)| (bucket.clone(), key.clone()))
976            .collect()
977    }
978
979    /// **The table `crates/model/src/image.rs` specifies**, row for row (`the_plan_table` there).
980    #[test]
981    fn the_plan_table() {
982        for (found, force, expected) in [
983            (Found::Absent, false, Plan::Build),
984            (Found::Absent, true, Plan::Build),
985            (Found::Ready, false, Plan::Reuse),
986            (Found::Ready, true, Plan::Delete),
987            (Found::Building, false, Plan::Wait),
988            (Found::Building, true, Plan::Wait),
989            (Found::Failed, false, Plan::Delete),
990            (Found::Failed, true, Plan::Delete),
991            (Found::Deleting, false, Plan::AwaitAbsent),
992            (Found::Deleting, true, Plan::AwaitAbsent),
993        ] {
994            assert_eq!(plan(found, force), expected, "{found:?} force={force}");
995        }
996        for (spelling, found) in [
997            ("CREATING", Found::Building),
998            ("UPDATING", Found::Building),
999            ("CREATED", Found::Ready),
1000            ("UPDATED", Found::Ready),
1001            ("ACTIVE", Found::Ready),
1002            ("CREATE_FAILED", Found::Failed),
1003            ("UPDATE_FAILED", Found::Failed),
1004            ("DELETE_FAILED", Found::Failed),
1005            ("DELETING", Found::Deleting),
1006            ("DELETED", Found::Absent),
1007        ] {
1008            assert_eq!(Found::from_state(spelling), found, "{spelling}");
1009        }
1010    }
1011
1012    /// **IMAGE-6, the name.** `<prefix>-<hash12>`, the prefix reduced to the characters the
1013    /// service's `ImageName` admits and cut so the whole fits its 64; a prefix with nothing
1014    /// left is refused.
1015    #[test]
1016    fn the_name_is_the_sanitized_prefix_and_twelve_hex() {
1017        let hash = "0123456789abcdef".repeat(4);
1018        assert_eq!(
1019            ensured_image_name("harbor/task:1", &hash).expect("legal"),
1020            "harbor-task-1-0123456789ab"
1021        );
1022        let long = ensured_image_name(&"p".repeat(100), &hash).expect("legal");
1023        assert_eq!(long.len(), 64, "IMAGE-6: {long}");
1024        assert!(long.ends_with("-0123456789ab"));
1025        super::super::require_valid_image_name(&long).expect("the service's pattern");
1026        for empty in ["", "///", "   "] {
1027            let error = ensured_image_name(empty, &hash).expect_err("nothing left");
1028            assert_eq!(error.kind(), ErrorKind::InvalidArg);
1029        }
1030    }
1031
1032    /// **IMAGE-6, the identity.** The base image and the size class are part of the name,
1033    /// because an image is created at one size on one base: two tasks with one Dockerfile and
1034    /// different sizes need two images.
1035    #[test]
1036    fn the_identity_covers_the_base_and_the_size() {
1037        let hash = "a".repeat(64);
1038        let base = BaseImage::al2023();
1039        let identity = image_identity_hash(&hash, &base, SizeClass::DEFAULT);
1040        assert_eq!(
1041            identity,
1042            image_identity_hash(&hash, &base, SizeClass::DEFAULT)
1043        );
1044        assert_eq!(identity.len(), 64);
1045        let other_size = SizeClass::from_baseline_mib(4096).expect("a class");
1046        assert_ne!(
1047            identity,
1048            image_identity_hash(&hash, &base, other_size),
1049            "IMAGE-6: size"
1050        );
1051        let other_base = BaseImage {
1052            name: "al2023-2".to_string(),
1053            ..BaseImage::al2023()
1054        };
1055        assert_ne!(
1056            identity,
1057            image_identity_hash(&hash, &other_base, SizeClass::DEFAULT)
1058        );
1059        assert_ne!(
1060            identity,
1061            image_identity_hash(&"b".repeat(64), &base, SizeClass::DEFAULT)
1062        );
1063    }
1064
1065    /// **IMAGE-6 with a pinned base (#258).** An unpinned identity is the historical one
1066    /// exactly, so every ensured image keeps its name; a pin is identity, and two pins differ.
1067    #[test]
1068    fn a_pinned_base_version_is_part_of_the_identity() {
1069        let hash = "a".repeat(64);
1070        let base = BaseImage::al2023();
1071        let unpinned = image_identity_hash(&hash, &base, SizeClass::DEFAULT);
1072        assert_eq!(
1073            pinned_identity_hash(&hash, &base, None, SizeClass::DEFAULT),
1074            unpinned,
1075            "no pin, the name every ensured image already has"
1076        );
1077        let one = pinned_identity_hash(&hash, &base, Some("1"), SizeClass::DEFAULT);
1078        assert_ne!(one, unpinned);
1079        assert_ne!(
1080            one,
1081            pinned_identity_hash(&hash, &base, Some("0"), SizeClass::DEFAULT)
1082        );
1083    }
1084
1085    /// **The create fields an ensure carries (#258).** Project files, a pinned base, logging,
1086    /// identity repair, `inherit_workdir` and the hook timeouts name a different image (#280),
1087    /// and all of them reach the create request a build sends.
1088    #[test]
1089    fn a_created_image_carries_every_field_of_the_request() {
1090        use crate::control::artifact::{Ecosystem, ProjectFiles};
1091
1092        let plane = crate::testing::control_plane(
1093            Arc::new(FakeControlPlane::new()),
1094            Region::UsEast1,
1095            Arc::new(TestClock::new()),
1096        );
1097        let plain = prepare(&plane, request()).expect("prepares");
1098
1099        let mut with_project = request();
1100        with_project.project_files = Some(ProjectFiles {
1101            ecosystem: Ecosystem::Uv,
1102            manifest: b"[project]\nname = \"demo\"\n".to_vec(),
1103            lockfile: b"version = 1\n".to_vec(),
1104        });
1105        // The wrapped Dockerfile doesn't COPY the pair, which core refuses (#74), so the
1106        // project's own default stanza stands in.
1107        with_project.dockerfile = crate::control::artifact::default_dockerfile(
1108            crate::control::DEFAULT_AGENT_PORT,
1109            None,
1110            &BaseImage::al2023(),
1111            Some(Ecosystem::Uv),
1112        );
1113        with_project.base_image = Some(BaseImage::al2023());
1114        let mut without_project = with_project.clone();
1115        without_project.project_files = None;
1116        let project = prepare(&plane, with_project).expect("prepares");
1117        assert!(project.create.project_files.is_some());
1118        assert_ne!(
1119            project.name,
1120            prepare(&plane, without_project).expect("prepares").name,
1121            "the project files enter the artifact, so they name a different image"
1122        );
1123
1124        let mut pinned = request();
1125        pinned.base_image_version = Some("1".to_string());
1126        let pinned = prepare(&plane, pinned).expect("prepares");
1127        assert_ne!(pinned.name, plain.name, "a pin names a different image");
1128        assert_eq!(pinned.create.base_image_version.as_deref(), Some("1"));
1129
1130        let mut configured = request();
1131        configured.log_group = Some("/team/builds".to_string());
1132        configured.log_stream = Some("task".to_string());
1133        configured.repair_guest_identity = true;
1134        configured.inherit_workdir = true;
1135        configured.run_hook_timeout =
1136            crate::hooks::RunHookTimeout::try_new(7).expect("a legal timeout");
1137        let configured = prepare(&plane, configured).expect("prepares");
1138        assert_ne!(
1139            configured.name, plain.name,
1140            "the create-only fields are identity (#280)"
1141        );
1142        assert_eq!(configured.create.log_group.as_deref(), Some("/team/builds"));
1143        assert_eq!(configured.create.log_stream.as_deref(), Some("task"));
1144        assert!(configured.create.repair_guest_identity);
1145        assert!(configured.create.inherit_workdir);
1146        assert_eq!(configured.create.run_hook_timeout.as_secs(), 7);
1147    }
1148
1149    /// **Each create-only field names a different image (#280).** The service sets identity
1150    /// repair, `inherit_workdir`, the hook timeouts and the log destination when it creates
1151    /// an image, and no later call changes them, so an image reused under a name that left
1152    /// them out can be one built without them. Each one changed alone renames the image, and
1153    /// the same values name the same image again.
1154    #[test]
1155    fn each_create_only_field_names_a_different_image() {
1156        let plane = crate::testing::control_plane(
1157            Arc::new(FakeControlPlane::new()),
1158            Region::UsEast1,
1159            Arc::new(TestClock::new()),
1160        );
1161        let plain = prepare(&plane, request()).expect("prepares").name;
1162        assert_eq!(
1163            prepare(&plane, request()).expect("prepares").name,
1164            plain,
1165            "equal inputs name one image"
1166        );
1167
1168        type Change = fn(&mut EnsureImageRequest);
1169        let changes: [(&str, Change); 7] = [
1170            ("repair_guest_identity", |r| r.repair_guest_identity = true),
1171            ("inherit_workdir", |r| r.inherit_workdir = true),
1172            ("run_hook_timeout", |r| {
1173                r.run_hook_timeout = crate::hooks::RunHookTimeout::try_new(7).expect("legal");
1174            }),
1175            ("build_hook_timeout", |r| {
1176                r.build_hook_timeout = crate::hooks::BuildHookTimeout::try_new(901).expect("legal");
1177            }),
1178            ("log_group", |r| {
1179                r.log_group = Some("/team/builds".to_string());
1180            }),
1181            ("another log_group", |r| {
1182                r.log_group = Some("/team/other".to_string());
1183            }),
1184            ("log_stream", |r| {
1185                r.log_group = Some("/team/builds".to_string());
1186                r.log_stream = Some("task".to_string());
1187            }),
1188        ];
1189        let mut names = vec![plain.clone()];
1190        for (field, change) in changes {
1191            let mut changed = request();
1192            change(&mut changed);
1193            let name = prepare(&plane, changed).expect("prepares").name;
1194            assert_ne!(name, plain, "{field} is create-only, so it is identity");
1195            assert!(
1196                !names.contains(&name),
1197                "{field} names an image another field already names: {name}"
1198            );
1199            names.push(name);
1200        }
1201
1202        // Tags are not: a reused image keeps its own.
1203        let mut tagged = request();
1204        tagged.tags.insert("team".to_string(), "x".to_string());
1205        assert_eq!(prepare(&plane, tagged).expect("prepares").name, plain);
1206    }
1207
1208    /// `from_create` hands an ensure every field a create request carries, and names the
1209    /// image the way an ensure of the same inputs would.
1210    #[test]
1211    fn an_ensure_from_a_create_request_keeps_its_fields() {
1212        let mut create = CreateImageRequest::new("task", aarch64_daemon(b"daemon"), "", ROLE);
1213        create.dockerfile = Some(dockerfile());
1214        create.base_image = BaseImage::from_dockerfile(&dockerfile()).expect("a FROM");
1215        create.base_image_version = Some("1".to_string());
1216        create.log_group = Some("/g".to_string());
1217        create.repair_guest_identity = true;
1218        create.tags.insert("team".to_string(), "x".to_string());
1219        let from = EnsureImageRequest::from_create(create, dockerfile(), BUCKET);
1220        assert_eq!(from.name_prefix, "task");
1221        assert_eq!(from.s3_bucket, BUCKET);
1222        assert_eq!(from.base_image_version.as_deref(), Some("1"));
1223        assert_eq!(from.log_group.as_deref(), Some("/g"));
1224        assert!(from.repair_guest_identity);
1225        assert_eq!(from.tags.get("team").map(String::as_str), Some("x"));
1226
1227        let mut direct = request();
1228        direct.base_image_version = Some("1".to_string());
1229        direct.log_group = Some("/g".to_string());
1230        direct.repair_guest_identity = true;
1231        direct.s3_key_prefix = None;
1232        let plane = crate::testing::control_plane(
1233            Arc::new(FakeControlPlane::new()),
1234            Region::UsEast1,
1235            Arc::new(TestClock::new()),
1236        );
1237        assert_eq!(
1238            prepare(&plane, from).expect("prepares").name,
1239            prepare(&plane, direct).expect("prepares").name
1240        );
1241    }
1242
1243    /// **A caller-named build uploads through the build services (#258).** The artifact goes to
1244    /// the bucket and key its URI names, once, with the account unasked; a URI that isn't one
1245    /// puts nothing.
1246    #[tokio::test]
1247    async fn a_named_builds_artifact_is_put_where_its_uri_says() {
1248        let fake = Arc::new(FakeControlPlane::new());
1249        let services = Arc::new(FakeServices::default());
1250        let mut sandbox = sandbox(&fake, &services);
1251        let mut create = CreateImageRequest::new(
1252            "named",
1253            aarch64_daemon(b"daemon"),
1254            "s3://artifact-bucket/builds/named.zip",
1255            ROLE,
1256        );
1257        create.dockerfile = Some(dockerfile());
1258        create.base_image = BaseImage::from_dockerfile(&dockerfile()).expect("a FROM");
1259        sandbox.upload_artifact_for(&create).await.expect("uploads");
1260        let puts = services.puts.lock().expect("not poisoned").clone();
1261        assert_eq!(puts.len(), 1);
1262        assert_eq!(
1263            (puts[0].0.as_str(), puts[0].1.as_str()),
1264            ("artifact-bucket", "builds/named.zip")
1265        );
1266        assert_eq!(
1267            puts[0].2,
1268            sandbox
1269                .build_artifact_for(&create)
1270                .expect("the same artifact")
1271        );
1272        assert_eq!(services.account_calls.load(Ordering::SeqCst), 0);
1273        assert!(fake.calls().is_empty(), "no control-plane call");
1274
1275        create.code_artifact_uri = "artifact-bucket/named.zip".to_string();
1276        let error = sandbox
1277            .upload_artifact_for(&create)
1278            .await
1279            .expect_err("not an S3 URI");
1280        assert_eq!(error.kind(), ErrorKind::InvalidArg, "{error}");
1281        assert_eq!(uploads(&services).len(), 1, "nothing more was put");
1282    }
1283
1284    /// An artifact URI is `s3://<bucket>/<key>` with a bucket S3 admits, or it is refused.
1285    #[test]
1286    fn an_artifact_uri_names_a_bucket_and_a_key() {
1287        assert_eq!(
1288            s3_location("s3://artifact-bucket/a/b.zip").expect("an S3 URI"),
1289            ("artifact-bucket", "a/b.zip")
1290        );
1291        for uri in [
1292            "https://artifact-bucket/a.zip",
1293            "s3://artifact-bucket",
1294            "s3://artifact-bucket/",
1295            "s3://Not_A_Bucket/a.zip",
1296        ] {
1297            let error = s3_location(uri).expect_err(uri);
1298            assert_eq!(error.kind(), ErrorKind::InvalidArg, "{uri}: {error}");
1299        }
1300    }
1301
1302    /// **IMAGE-8, the key.** `<prefix>/<name>/artifact.zip`, with the prefix's slashes
1303    /// trimmed so no key carries an empty segment.
1304    #[test]
1305    fn the_artifact_key_is_content_addressed() {
1306        assert_eq!(
1307            artifact_key(Some("harbor/images"), "task-0123456789ab"),
1308            "harbor/images/task-0123456789ab/artifact.zip"
1309        );
1310        assert_eq!(artifact_key(Some("/harbor/"), "n"), "harbor/n/artifact.zip");
1311        assert_eq!(artifact_key(None, "n"), "n/artifact.zip");
1312        assert_eq!(artifact_key(Some(""), "n"), "n/artifact.zip");
1313    }
1314
1315    /// **IMAGE-9, reuse.** A ready image is returned as it is: one account lookup, one
1316    /// describe, no upload, no create.
1317    #[tokio::test]
1318    async fn a_ready_image_is_reused_with_no_upload_and_no_create() {
1319        let name = expected_name();
1320        let fake = Arc::new(FakeControlPlane::new());
1321        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)));
1322        let services = Arc::new(FakeServices::default());
1323        let ensured = sandbox(&fake, &services)
1324            .ensure_image(request())
1325            .await
1326            .expect("reused");
1327        assert!(ensured.reused, "IMAGE-9");
1328        assert!(!ensured.uploaded);
1329        assert_eq!(ensured.image.name, name);
1330        assert_eq!(ensured.image.version, "1");
1331        assert_eq!(fake.operations(), ["GetMicrovmImage"], "IMAGE-9");
1332        assert!(uploads(&services).is_empty());
1333        assert!(
1334            fake.paths()[0].contains(&paths_encoded(&format!(
1335                "arn:aws:lambda:us-east-1:{ACCOUNT}:microvm-image:{name}"
1336            ))),
1337            "IMAGE-8: the describe names the ARN built from the account: {:?}",
1338            fake.paths()
1339        );
1340    }
1341
1342    fn paths_encoded(arn: &str) -> String {
1343        crate::control::transport::paths::encode_segment(arn)
1344    }
1345
1346    /// **IMAGE-9, wait.** A build already running under the name is waited out and reused;
1347    /// this caller uploads and creates nothing.
1348    #[tokio::test]
1349    async fn a_running_build_is_waited_out_and_reused() {
1350        let name = expected_name();
1351        let fake = Arc::new(FakeControlPlane::new());
1352        fake.answer("GetMicrovmImage", state(&name, "CREATING"))
1353            .answer("GetMicrovmImage", Answer::ok(ready(&name)));
1354        let services = Arc::new(FakeServices::default());
1355        let ensured = sandbox(&fake, &services)
1356            .ensure_image(request())
1357            .await
1358            .expect("waited");
1359        assert!(ensured.reused, "IMAGE-9");
1360        assert_eq!(fake.call_count("CreateMicrovmImage"), 0, "IMAGE-9");
1361        assert!(uploads(&services).is_empty(), "IMAGE-9");
1362    }
1363
1364    /// **IMAGE-8, the build.** No image under the name: the artifact goes to the
1365    /// content-addressed key, the create names that URI and the derived name, and the image
1366    /// this call built is not reused.
1367    #[tokio::test]
1368    async fn an_absent_image_is_uploaded_to_its_key_and_built() {
1369        let name = expected_name();
1370        let fake = Arc::new(FakeControlPlane::new());
1371        fake.answer("GetMicrovmImage", absent())
1372            .answer("GetMicrovmImage", state(&name, "CREATING"))
1373            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1374            .answer(
1375                "CreateMicrovmImage",
1376                Answer::created(fake::create_image_response(&name)),
1377            );
1378        let services = Arc::new(FakeServices::default());
1379        let mut request = request();
1380        request
1381            .tags
1382            .insert("harbor:task".to_string(), "t".to_string());
1383        let ensured = sandbox(&fake, &services)
1384            .ensure_image(request)
1385            .await
1386            .expect("built");
1387        assert!(!ensured.reused);
1388        assert!(ensured.uploaded);
1389        let key = format!("harbor/images/{name}/artifact.zip");
1390        assert_eq!(
1391            uploads(&services),
1392            [(BUCKET.to_string(), key.clone())],
1393            "IMAGE-8"
1394        );
1395        assert_eq!(ensured.artifact_uri, format!("s3://{BUCKET}/{key}"));
1396        let body = fake.first_body("CreateMicrovmImage");
1397        assert_eq!(body["name"], name.as_str());
1398        assert_eq!(body["codeArtifact"]["uri"], ensured.artifact_uri.as_str());
1399        assert_eq!(body["tags"]["harbor:task"], "t");
1400        assert_eq!(
1401            body["baseImageArn"], "arn:aws:lambda:us-east-1:aws:microvm-image:al2023-1",
1402            "the derived base keeps the managed name"
1403        );
1404        assert_eq!(services.account_calls.load(Ordering::SeqCst), 1);
1405
1406        // The uploaded bytes are the artifact the request describes.
1407        let (_, _, bytes) = services.puts.lock().expect("not poisoned")[0].clone();
1408        let archive = zip::ZipArchive::new(std::io::Cursor::new(bytes)).expect("a zip");
1409        assert_eq!(archive.len(), 2);
1410    }
1411
1412    /// **IMAGE-10, a failure.** A failed image is deleted, the name awaited free, and then
1413    /// rebuilt; the delete comes before the create.
1414    #[tokio::test]
1415    async fn a_failed_image_is_deleted_then_rebuilt() {
1416        let name = expected_name();
1417        let fake = Arc::new(FakeControlPlane::new());
1418        fake.answer("GetMicrovmImage", state(&name, "CREATE_FAILED"))
1419            .answer("GetMicrovmImage", absent())
1420            .answer("GetMicrovmImage", state(&name, "CREATING"))
1421            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1422            .answer(
1423                "ListMicrovmImageVersions",
1424                Answer::ok(fake::list_versions_response("1")),
1425            )
1426            .answer(
1427                "DeleteMicrovmImage",
1428                Answer::ok(fake::delete_image_response()),
1429            )
1430            .answer(
1431                "CreateMicrovmImage",
1432                Answer::created(fake::create_image_response(&name)),
1433            );
1434        let services = Arc::new(FakeServices::default());
1435        let ensured = sandbox(&fake, &services)
1436            .ensure_image(request())
1437            .await
1438            .expect("rebuilt");
1439        assert!(!ensured.reused);
1440        let operations = fake.operations();
1441        let deleted = operations
1442            .iter()
1443            .position(|op| *op == "DeleteMicrovmImage")
1444            .expect("IMAGE-10: deleted");
1445        let created = operations
1446            .iter()
1447            .position(|op| *op == "CreateMicrovmImage")
1448            .expect("rebuilt");
1449        assert!(deleted < created, "IMAGE-10: {operations:?}");
1450    }
1451
1452    /// **IMAGE-10, force.** A ready image is rebuilt only when the caller forces it.
1453    #[tokio::test]
1454    async fn a_ready_image_is_rebuilt_only_under_force() {
1455        let name = expected_name();
1456        let fake = Arc::new(FakeControlPlane::new());
1457        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)))
1458            .answer("GetMicrovmImage", absent())
1459            .answer("GetMicrovmImage", state(&name, "CREATING"))
1460            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1461            .answer(
1462                "ListMicrovmImageVersions",
1463                Answer::ok(fake::list_versions_response("1")),
1464            )
1465            .answer(
1466                "DeleteMicrovmImage",
1467                Answer::ok(fake::delete_image_response()),
1468            )
1469            .answer(
1470                "CreateMicrovmImage",
1471                Answer::created(fake::create_image_response(&name)),
1472            );
1473        let services = Arc::new(FakeServices::default());
1474        let mut forced = request();
1475        forced.force = true;
1476        let ensured = sandbox(&fake, &services)
1477            .ensure_image(forced)
1478            .await
1479            .expect("rebuilt");
1480        assert!(!ensured.reused, "IMAGE-10");
1481        assert_eq!(fake.call_count("DeleteMicrovmImage"), 1, "IMAGE-10");
1482        assert_eq!(fake.call_count("CreateMicrovmImage"), 1);
1483    }
1484
1485    /// **IMAGE-10, a deletion under way.** The name is awaited free, then built.
1486    #[tokio::test]
1487    async fn a_deletion_under_way_is_awaited_then_built() {
1488        let name = expected_name();
1489        let fake = Arc::new(FakeControlPlane::new());
1490        fake.answer("GetMicrovmImage", state(&name, "DELETING"))
1491            .answer("GetMicrovmImage", absent())
1492            .answer("GetMicrovmImage", state(&name, "CREATING"))
1493            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1494            .answer(
1495                "CreateMicrovmImage",
1496                Answer::created(fake::create_image_response(&name)),
1497            );
1498        let services = Arc::new(FakeServices::default());
1499        let ensured = sandbox(&fake, &services)
1500            .ensure_image(request())
1501            .await
1502            .expect("built");
1503        assert!(!ensured.reused);
1504        assert_eq!(
1505            fake.call_count("DeleteMicrovmImage"),
1506            0,
1507            "not this caller's delete"
1508        );
1509    }
1510
1511    /// **IMAGE-11, the race.** The create is refused because a sibling created the name
1512    /// first: describe again, wait for the sibling's build, and return it as reused.
1513    #[tokio::test]
1514    async fn a_refused_create_joins_the_winners_build() {
1515        let name = expected_name();
1516        let fake = Arc::new(FakeControlPlane::new());
1517        fake.answer("GetMicrovmImage", absent())
1518            .answer("GetMicrovmImage", state(&name, "CREATING"))
1519            .answer("GetMicrovmImage", state(&name, "CREATING"))
1520            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1521            .answer(
1522                "CreateMicrovmImage",
1523                Answer::failure(409, "An image with this name already exists"),
1524            );
1525        let services = Arc::new(FakeServices::default());
1526        let ensured = sandbox(&fake, &services)
1527            .ensure_image(request())
1528            .await
1529            .expect("IMAGE-11: joined");
1530        assert!(
1531            ensured.reused,
1532            "IMAGE-11: this caller's create did not build it"
1533        );
1534        assert_eq!(
1535            ensured.image.state, "CREATED",
1536            "IMAGE-11: returned only once ready"
1537        );
1538        assert_eq!(fake.call_count("CreateMicrovmImage"), 1);
1539    }
1540
1541    /// **IMAGE-11, not a race.** A refused create with nothing under the name afterwards is
1542    /// the create's own failure, surfaced as it came.
1543    #[tokio::test]
1544    async fn a_refused_create_with_no_winner_surfaces_the_refusal() {
1545        let fake = Arc::new(FakeControlPlane::new());
1546        fake.answer("GetMicrovmImage", absent()).answer(
1547            "CreateMicrovmImage",
1548            Answer::failure(400, "buildRoleArn is invalid"),
1549        );
1550        let services = Arc::new(FakeServices::default());
1551        let error = sandbox(&fake, &services)
1552            .ensure_image(request())
1553            .await
1554            .expect_err("no winner");
1555        assert!(
1556            error.to_string().contains("buildRoleArn is invalid"),
1557            "{error}"
1558        );
1559    }
1560
1561    /// **IMAGE-8, once per sandbox.** Two ensures on one sandbox make one account lookup.
1562    #[tokio::test]
1563    async fn the_account_is_resolved_once_per_sandbox() {
1564        let name = expected_name();
1565        let fake = Arc::new(FakeControlPlane::new());
1566        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)));
1567        let services = Arc::new(FakeServices::default());
1568        let mut sandbox = sandbox(&fake, &services);
1569        sandbox.ensure_image(request()).await.expect("first");
1570        sandbox.ensure_image(request()).await.expect("second");
1571        assert_eq!(services.account_calls.load(Ordering::SeqCst), 1, "IMAGE-8");
1572    }
1573
1574    /// The account lookup of the test below: it records whether the other task's timer had
1575    /// fired by the time the ensure asked, then refuses, so the ensure goes no further.
1576    struct TimedLookup {
1577        fired: Arc<std::sync::atomic::AtomicBool>,
1578        seen: Mutex<Option<bool>>,
1579    }
1580
1581    impl BuildServices for TimedLookup {
1582        fn caller_account(&self) -> futures_util::future::BoxFuture<'_, Result<String, Error>> {
1583            *self.seen.lock().expect("not poisoned") = Some(self.fired.load(Ordering::SeqCst));
1584            Box::pin(async { Err(Error::new(ErrorKind::Credentials, "the probe's refusal")) })
1585        }
1586
1587        fn put_object<'a>(
1588            &'a self,
1589            _bucket: &'a str,
1590            _key: &'a str,
1591            _bytes: Vec<u8>,
1592        ) -> futures_util::future::BoxFuture<'a, Result<(), Error>> {
1593            Box::pin(async { panic!("the probe refuses before any upload") })
1594        }
1595    }
1596
1597    /// **`ensure_image` prepares its artifact off the caller's worker (#309).** On a
1598    /// current-thread runtime, a one-millisecond timer another task set before the ensure
1599    /// began fires while a daemon-sized artifact is hashed and zipped, before the ensure asks
1600    /// for its account.
1601    ///
1602    /// Inline, `prepare` holds the runtime's only worker from its first line to that lookup,
1603    /// so no other task runs in between: a concurrent ensure's credential fetch parked
1604    /// behind it is what timed out under instance-role credentials. The daemon here is 2 MiB
1605    /// that deflate can't shrink, like the real one, which takes hundreds of milliseconds to
1606    /// prepare in a debug build, so the timer has that long to fire.
1607    ///
1608    /// **Falsification**: call `prepare` inline in `Sandbox::ensure_image` again and the
1609    /// lookup finds the timer unfired.
1610    #[tokio::test(flavor = "current_thread")]
1611    async fn the_artifact_is_prepared_while_the_callers_other_tasks_run() {
1612        let fired = Arc::new(std::sync::atomic::AtomicBool::new(false));
1613        let lookup = Arc::new(TimedLookup {
1614            fired: Arc::clone(&fired),
1615            seen: Mutex::new(None),
1616        });
1617        let fake = Arc::new(FakeControlPlane::new());
1618        let plane = crate::testing::control_plane(
1619            fake.clone(),
1620            Region::UsEast1,
1621            Arc::new(TestClock::new()),
1622        );
1623        let mut sandbox = Sandbox::with_control_plane(plane).with_build_services(lookup.clone());
1624        let mut request = request();
1625        // xorshift64: bytes with no repeats for deflate to find, from no dependency, behind
1626        // the aarch64 header the preflight asks for.
1627        let mut state = 0x9e37_79b9_7f4a_7c15_u64;
1628        let payload: Vec<u8> = (0..(2 << 20) / 8)
1629            .flat_map(|_| {
1630                state ^= state << 13;
1631                state ^= state >> 7;
1632                state ^= state << 17;
1633                state.to_le_bytes()
1634            })
1635            .collect();
1636        request.binary = aarch64_daemon(&payload);
1637
1638        let timer = tokio::spawn(async move {
1639            tokio::time::sleep(Duration::from_millis(1)).await;
1640            fired.store(true, Ordering::SeqCst);
1641        });
1642        let error = sandbox
1643            .ensure_image(request)
1644            .await
1645            .expect_err("the probe refuses the lookup");
1646        assert!(error.to_string().contains("the probe's refusal"), "{error}");
1647        assert_eq!(
1648            *lookup.seen.lock().expect("not poisoned"),
1649            Some(true),
1650            "the other task's timer fires while the artifact is prepared, not after"
1651        );
1652        assert!(fake.calls().is_empty(), "no control-plane call came first");
1653        timer.await.expect("the timer task");
1654    }
1655
1656    /// Every local refusal costs zero calls — not the account lookup, not a describe, not
1657    /// the upload.
1658    #[tokio::test]
1659    async fn a_locally_refused_request_makes_no_call() {
1660        let fake = Arc::new(FakeControlPlane::new());
1661        let services = Arc::new(FakeServices::default());
1662        let mut sandbox = sandbox(&fake, &services);
1663        for (mutate, cause) in [
1664            (
1665                Box::new(|r: &mut EnsureImageRequest| {
1666                    r.dockerfile = "FROM x\nENTRYPOINT [\"/bin/sh\"]\nCMD [\"/agentd\"]\n".into()
1667                }) as Box<dyn Fn(&mut EnsureImageRequest)>,
1668                "ENTRYPOINT",
1669            ),
1670            (Box::new(|r| r.name_prefix = "//".into()), "prefix"),
1671            (Box::new(|r| r.s3_bucket = String::new()), "bucket"),
1672            (
1673                Box::new(|r| r.build_role_arn = "not-an-arn".into()),
1674                "buildRoleArn",
1675            ),
1676            (
1677                Box::new(|r| r.s3_key_prefix = Some("a\nb".into())),
1678                "key prefix",
1679            ),
1680        ] {
1681            let mut request = request();
1682            mutate(&mut request);
1683            let error = sandbox.ensure_image(request).await.expect_err(cause);
1684            assert_eq!(error.kind(), ErrorKind::InvalidArg, "{error}");
1685            assert!(error.to_string().contains(cause), "{cause}: {error}");
1686        }
1687        assert!(fake.calls().is_empty());
1688        assert_eq!(services.account_calls.load(Ordering::SeqCst), 0);
1689        assert!(uploads(&services).is_empty());
1690    }
1691
1692    /// The build wait honours `wait_timeout`: a build that never settles is a timeout at the
1693    /// caller's deadline, not the 45-minute default.
1694    #[tokio::test]
1695    async fn the_wait_timeout_bounds_the_build_wait() {
1696        let name = expected_name();
1697        let fake = Arc::new(FakeControlPlane::new());
1698        fake.answer("GetMicrovmImage", absent())
1699            .answer("GetMicrovmImage", state(&name, "CREATING"))
1700            .answer(
1701                "CreateMicrovmImage",
1702                Answer::created(fake::create_image_response(&name)),
1703            )
1704            .answer(
1705                "ListMicrovmImageBuilds",
1706                Answer::ok(fake::list_builds_response("IN_PROGRESS")),
1707            )
1708            .answer(
1709                "ListMicrovmImageVersions",
1710                Answer::ok(fake::list_versions_response("1")),
1711            );
1712        let services = Arc::new(FakeServices::default());
1713        let mut request = request();
1714        request.wait_timeout = Some(Duration::from_secs(60));
1715        let error = sandbox(&fake, &services)
1716            .ensure_image(request)
1717            .await
1718            .expect_err("never settles");
1719        assert_eq!(error.kind(), ErrorKind::Timeout, "{error}");
1720        assert!(
1721            fake.call_count("GetMicrovmImage") < 10,
1722            "bounded by 60s, not 45 minutes"
1723        );
1724    }
1725}