Create ImageRef type to reference local images

pipex · pipex · commit 5501abc3dd18 · 2025-09-30T11:37:16.000-03:00
Local images may be referenced by URI or by Id. This reference is used
when reading service information from containers and mapping a container
to the backend image URI

Change-type: patch
diff --git a/helios-oci/src/lib.rs b/helios-oci/src/lib.rs
@@ -12,7 +12,7 @@ mod container;
 pub use container::{Container, LocalContainer};
 
 mod models;
-pub use models::{ImageUri, InvalidImageUriError};
+pub use models::{ImageRef, ImageUri, InvalidImageUriError};
 
 mod registry;
 pub use registry::{RegistryAuth, RegistryAuthClient, RegistryAuthError};
diff --git a/helios-oci/src/models.rs b/helios-oci/src/models.rs
@@ -178,6 +178,45 @@ impl<'de> Deserialize<'de> for ImageUri {
     }
 }
 
+/// An image reference is either a image URI or a content addressable image ID
+#[derive(Serialize, Debug, Clone, PartialEq, Eq)]
+#[serde(untagged)]
+pub enum ImageRef {
+    /// A URI reference
+    Uri(ImageUri),
+
+    /// A content addressable image id
+    Id(String),
+}
+
+impl ImageRef {
+    /// Convenience method to get the digest of an image ref
+    ///
+    /// Returns None if the image is not a Uri or the Uri does not have a digest
+    pub fn digest(&self) -> &Option<String> {
+        if let Self::Uri(uri) = self {
+            uri.digest()
+        } else {
+            &None
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for ImageRef {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let s: String = String::deserialize(deserializer)?;
+        if s.starts_with("sha256:") {
+            return Ok(ImageRef::Id(s));
+        }
+
+        let uri: ImageUri = s.parse().map_err(serde::de::Error::custom)?;
+        Ok(ImageRef::Uri(uri))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/helios-state/src/models/service.rs b/helios-state/src/models/service.rs
@@ -5,7 +5,7 @@ use std::str::FromStr;
 
 use serde::{Deserialize, Serialize};
 
-use crate::oci::ImageUri;
+use crate::oci::ImageRef;
 use crate::util::types::Uuid;
 
 // We don't want to fail if the service is supervised but it doesn't have an app-uuid,
@@ -92,7 +92,7 @@ impl FromStr for ServiceContainerName {
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct Service {
     pub id: u32,
-    pub image: ImageUri,
+    pub image: ImageRef,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
@@ -101,8 +101,8 @@ pub struct TargetService {
     #[serde(default)]
     pub id: u32,
 
-    /// Service image URI
-    pub image: ImageUri,
+    /// Service image
+    pub image: ImageRef,
 }
 
 impl From<Service> for TargetService {
diff --git a/helios-state/src/read.rs b/helios-state/src/read.rs
@@ -4,7 +4,7 @@ use thiserror::Error;
 use tracing::instrument;
 
 use crate::oci::{
-    Client as Docker, Error as DockerError, ImageUri, InvalidImageUriError, WithContext,
+    Client as Docker, Error as DockerError, ImageRef, ImageUri, InvalidImageUriError, WithContext,
 };
 use crate::util::state;
 use crate::util::types::{OperatingSystem, Uuid};
@@ -99,18 +99,19 @@ pub async fn read(
 
         // Insert the service and link it to the image if there is state
         // metadata about the image
-        // FIXME: the in-memory service should exist if the container exists so
-        // we need to use an image reference if the state variable does not exist
-        if let Some(image) = svc_img {
-            let svc_id: u32 = labels
-                .get("io.balena.service-id")
-                .and_then(|id| id.parse().ok())
-                .unwrap_or(0);
-
-            release
-                .services
-                .insert(service_name, Service { id: svc_id, image });
-        }
+        let image = if let Some(svc_img) = svc_img {
+            ImageRef::Uri(svc_img)
+        } else {
+            ImageRef::Id(local_container.image_id)
+        };
+
+        let svc_id: u32 = labels
+            .get("io.balena.service-id")
+            .and_then(|id| id.parse().ok())
+            .unwrap_or(0);
+        release
+            .services
+            .insert(service_name, Service { id: svc_id, image });
     }
 
     Ok(device)
diff --git a/helios-state/src/worker.rs b/helios-state/src/worker.rs
@@ -11,7 +11,7 @@ use mahler::{
 use tokio::sync::RwLock;
 
 use super::util::state;
-use crate::oci::{Client as Docker, Error as DockerError, ImageUri};
+use crate::oci::{Client as Docker, Error as DockerError, ImageRef, ImageUri};
 use crate::oci::{Credentials, RegistryAuth, RegistryAuthClient, RegistryAuthError};
 use crate::util::types::Uuid;
 
@@ -209,23 +209,34 @@ fn request_registry_token_for_new_images(
                 release
                     .services
                     .iter()
-                    .filter(|(service_name, svc)| {
-                        // if the service image has not been downloaded
-                        !device.images.contains_key(&svc.image) &&
-                        // and there is no service from a different release
+                    .filter_map(|(svc_name, svc)| {
+                        if let ImageRef::Uri(img_uri) = &svc.image {
+                            Some((svc_name, img_uri))
+                        } else {
+                            None
+                        }
+                    })
+                    .filter(|(svc_name, svc_img)| {
+                        // ignore the image if it already exists
+                        if device.images.contains_key(svc_img) {
+                            return false;
+                        }
+
                         find_installed_service(
                             &device,
                             app_uuid.clone(),
                             commit.clone(),
-                            (*service_name).clone(),
+                            (*svc_name).clone(),
                         )
-                        // or there is a service and the digest doesn't match the target digest
+                        // if the image is for a new service or the existing service has a
+                        // different digest (which means a new download), then add it to the
+                        // authorization list
                         .is_none_or(|s| {
-                            s.image.digest().is_none() || s.image.digest() != svc.image.digest()
+                            s.image.digest().is_none() || s.image.digest() != svc_img.digest()
                         })
                     })
                     // then select the image
-                    .map(|(_, svc)| &svc.image)
+                    .map(|(_, svc_img)| svc_img)
             })
         })
         .collect();
@@ -265,29 +276,36 @@ fn fetch_release_images(
     let images_to_install: Vec<ImageUri> = tgt_release
         .services
         .into_iter()
-        .filter(|(service_name, svc)| {
-            // if the service image has not been downloaded
-            !device.images.contains_key(&svc.image) &&
-            // and there is no service from a different release
+        .filter_map(|(svc_name, svc)| {
+            if let ImageRef::Uri(img_uri) = svc.image {
+                Some((svc_name, img_uri))
+            } else {
+                None
+            }
+        })
+        .filter(|(svc_name, svc_img)| {
+            // ignore the image if it already exists
+            if device.images.contains_key(svc_img) {
+                return false;
+            }
+
             find_installed_service(
-                    &device,
-                    app_uuid.clone(),
-                    commit.clone(),
-                    (*service_name).clone(),
-                )
-                // or if there is a service, then only chose it if it has the same digest
-                // as the target service
-                .is_none_or(|s| {
-                    s.image.digest().is_some() && svc.image.digest() == s.image.digest()
-                })
+                &device,
+                app_uuid.clone(),
+                commit.clone(),
+                (*svc_name).clone(),
+            )
+            // select the image if it is for a new service or the existing service image has the
+            // same digest (which means we are just re-tagging)
+            .is_none_or(|s| s.image.digest().is_some() && svc_img.digest() == s.image.digest())
         })
         // remove duplicate digests
-        .fold(Vec::<ImageUri>::new(), |mut acc, (_, svc)| {
+        .fold(Vec::<ImageUri>::new(), |mut acc, (_, svc_img)| {
             if acc
                 .iter()
-                .all(|img| img.digest().is_none() || img.digest() != svc.image.digest())
+                .all(|img| img.digest().is_none() || img.digest() != svc_img.digest())
             {
-                acc.push(svc.image);
+                acc.push(svc_img);
             }
             acc
         });
@@ -415,18 +433,24 @@ fn fetch_service_image(
     Target(tgt): Target<TargetService>,
     Args((app_uuid, commit, service_name)): Args<(Uuid, Uuid, String)>,
 ) -> Option<Task> {
-    // tell the planner to skip this task if the image already exists
-    if device.images.contains_key(&tgt.image) {
+    let tgt_img = if let ImageRef::Uri(img) = tgt.image {
+        // Skip this task if the image already exists
+        if device.images.contains_key(&img) {
+            return None;
+        }
+        img
+    } else {
+        // also skip if the target image is not a Uri
         return None;
-    }
+    };
 
     // If there is a service with the same name in any other releases of the service
     // and the service image has as different digest then we skip the fetch as
     // that pull will need deltas and needs to be handled by another task
     if find_installed_service(&device, app_uuid, commit, service_name)
-        .is_none_or(|svc| svc.image.digest().is_some() && tgt.image.digest() == svc.image.digest())
+        .is_none_or(|svc| svc.image.digest().is_some() && tgt_img.digest() == svc.image.digest())
     {
-        return Some(create_image.with_arg("image_name", tgt.image));
+        return Some(create_image.with_arg("image_name", tgt_img));
     }
 
     None
@@ -443,7 +467,11 @@ fn install_service(
     Path(job_path): Path,
 ) -> IO<Option<Service>, state::ReadWriteError> {
     // Skip the task if the image for the service doesn't exist yet
-    if device.images.keys().all(|img| img != &tgt.image) {
+    if let ImageRef::Uri(tgt_img) = &tgt.image {
+        if !device.images.contains_key(tgt_img) {
+            return svc.into();
+        }
+    } else {
         return svc.into();
     }
 
@@ -472,10 +500,13 @@ fn remove_image(
     // only remove the image if it not being used by any service
     if device.apps.values().any(|app| {
         app.releases.values().any(|release| {
-            release
-                .services
-                .values()
-                .any(|s| s.image == image_name.clone())
+            release.services.values().any(|s| {
+                if let ImageRef::Uri(s_img) = &s.image {
+                    s_img == &image_name
+                } else {
+                    false
+                }
+            })
         })
     }) {
         return img_ptr.into();
