feat(vmm): OCI registry image discovery and pull

vmm/Cargo.toml

@@ -55,6 +55,7 @@ fatfs.workspace = true
 fscommon.workspace = true
 or-panic.workspace = true
 url.workspace = true
+reqwest.workspace = true
 
 [dev-dependencies]
 insta.workspace = true

vmm/rpc/proto/vmm_rpc.proto

@@ -350,6 +350,11 @@ service Vmm {
   rpc SvStop(Id) returns (google.protobuf.Empty);
   // Remove a stopped supervisor process by ID.
   rpc SvRemove(Id) returns (google.protobuf.Empty);
+
+  // List images available in the configured OCI registry.
+  rpc ListRegistryImages(google.protobuf.Empty) returns (RegistryImageListResponse);
+  // Pull an image from the OCI registry to local storage.
+  rpc PullRegistryImage(PullRegistryImageRequest) returns (google.protobuf.Empty);
 }
 
 // DHCP lease event reported by the host DHCP server.
@@ -365,6 +370,27 @@ message SvListResponse {
   repeated SvProcessInfo processes = 1;
 }
 
+// Available images discovered from the OCI registry.
+message RegistryImageListResponse {
+  repeated RegistryImageInfo images = 1;
+}
+
+// Metadata for an image tag in the OCI registry.
+message RegistryImageInfo {
+  // Tag name (e.g., "0.5.8", "nvidia-0.5.8")
+  string tag = 1;
+  // Whether this image is already downloaded locally
+  bool local = 2;
+  // Whether this image is currently being pulled
+  bool pulling = 3;
+}
+
+// Request to pull an image from the OCI registry.
+message PullRegistryImageRequest {
+  // Tag to pull (e.g., "0.5.8")
+  string tag = 1;
+}
+
 // Information about a single supervisor process.
 message SvProcessInfo {
   string id = 1;

vmm/src/app.rs

@@ -35,6 +35,7 @@ pub use qemu::{VmConfig, VmWorkDir};
 mod id_pool;
 mod image;
 mod qemu;
+pub(crate) mod registry;
 
 #[derive(Deserialize, Serialize, Debug, Clone)]
 pub struct PortMapping {
@@ -124,6 +125,8 @@ pub struct App {
     pub supervisor: SupervisorClient,
     state: Arc<Mutex<AppState>>,
     forward_service: Arc<tokio::sync::Mutex<ForwardService>>,
+    /// Tags currently being pulled from the image registry.
+    pub(crate) pulling_tags: Arc<Mutex<std::collections::HashSet<String>>>,
 }
 
 impl App {
@@ -152,6 +155,7 @@ impl App {
             })),
             config: Arc::new(config),
             forward_service: Arc::new(tokio::sync::Mutex::new(ForwardService::new())),
+            pulling_tags: Arc::new(Mutex::new(std::collections::HashSet::new())),
         }
     }
 

vmm/src/app/registry.rs

@@ -0,0 +1,287 @@
+// SPDX-FileCopyrightText: © 2025 Phala Network <dstack@phala.network>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::Path;
+use std::process::Stdio;
+
+use anyhow::{bail, Context, Result};
+use serde::Deserialize;
+use tokio::process::Command;
+use tracing::info;
+
+/// List tags from a Docker Registry HTTP API v2 endpoint.
+///
+/// `image_ref` is in the form `registry.example.com/repo/name`.
+pub async fn list_registry_tags(image_ref: &str) -> Result<Vec<String>> {
+    let (registry, repo) = parse_image_ref(image_ref)?;
+
+    // Try Docker Registry HTTP API v2
+    let url = format!("https://{registry}/v2/{repo}/tags/list");
+    info!("fetching registry tags from {url}");
+
+    let client = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(30))
+        .build()?;
+
+    let response = client
+        .get(&url)
+        .send()
+        .await
+        .context("failed to fetch registry tags")?;
+
+    if response.status() == reqwest::StatusCode::UNAUTHORIZED {
+        // Try with anonymous token auth (Docker Hub style)
+        return list_tags_with_token(&client, &registry, &repo).await;
+    }
+
+    if !response.status().is_success() {
+        bail!(
+            "registry returned HTTP {}: {}",
+            response.status(),
+            response.text().await.unwrap_or_default()
+        );
+    }
+
+    let tag_list: TagList = response
+        .json()
+        .await
+        .context("failed to parse registry tag list")?;
+
+    Ok(tag_list.tags.unwrap_or_default())
+}
+
+/// Handle token-based auth (Docker Hub / registries requiring Bearer token).
+async fn list_tags_with_token(
+    client: &reqwest::Client,
+    registry: &str,
+    repo: &str,
+) -> Result<Vec<String>> {
+    // Fetch token from the registry's token endpoint
+    let token_url =
+        format!("https://{registry}/v2/token?service={registry}&scope=repository:{repo}:pull");
+    let token_resp = client.get(&token_url).send().await;
+
+    // If the token endpoint doesn't exist, try the standard Docker Hub approach
+    let token = match token_resp {
+        Ok(resp) if resp.status().is_success() => {
+            let token_data: TokenResponse = resp.json().await?;
+            token_data.token
+        }
+        _ => {
+            bail!("registry requires authentication but token exchange failed");
+        }
+    };
+
+    let url = format!("https://{registry}/v2/{repo}/tags/list");
+    let response = client
+        .get(&url)
+        .bearer_auth(&token)
+        .send()
+        .await
+        .context("failed to fetch registry tags with token")?;
+
+    if !response.status().is_success() {
+        bail!(
+            "registry returned HTTP {} after auth: {}",
+            response.status(),
+            response.text().await.unwrap_or_default()
+        );
+    }
+
+    let tag_list: TagList = response
+        .json()
+        .await
+        .context("failed to parse registry tag list")?;
+
+    Ok(tag_list.tags.unwrap_or_default())
+}
+
+/// Pull an image from registry and extract to the local image directory.
+///
+/// Uses `docker pull` + `docker create` + `docker export` to extract files.
+pub async fn pull_and_extract(image_ref: &str, tag: &str, image_path: &Path) -> Result<()> {
+    let full_ref = format!("{image_ref}:{tag}");
+    info!("pulling image {full_ref}");
+
+    // docker pull
+    let output = Command::new("docker")
+        .args(["pull", &full_ref])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .output()
+        .await
+        .context("failed to execute docker pull")?;
+
+    if !output.status.success() {
+        bail!(
+            "docker pull failed: {}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    // Determine output directory name from image metadata
+    let output_dir = determine_output_dir(image_ref, tag, image_path).await?;
+    if output_dir.exists() {
+        bail!("image directory already exists: {}", output_dir.display());
+    }
+
+    // Create temp dir, extract, then rename
+    let tmp_dir = image_path.join(format!(".tmp-pull-{tag}"));
+    if tmp_dir.exists() {
+        fs_err::remove_dir_all(&tmp_dir).context("failed to clean up temp dir")?;
+    }
+    fs_err::create_dir_all(&tmp_dir)?;
+
+    // docker create (don't start)
+    let output = Command::new("docker")
+        .args(["create", &full_ref, "/nonexistent"])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .output()
+        .await
+        .context("failed to create container")?;
+
+    if !output.status.success() {
+        let _ = fs_err::remove_dir_all(&tmp_dir);
+        bail!(
+            "docker create failed: {}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    let container_id = String::from_utf8_lossy(&output.stdout).trim().to_string();
+
+    // docker export | tar extract
+    let result = extract_container(&container_id, &tmp_dir).await;
+
+    // Always clean up container
+    let _ = Command::new("docker")
+        .args(["rm", &container_id])
+        .output()
+        .await;
+
+    result?;
+
+    // Remove docker artifact directories (FROM scratch creates these)
+    for dir in &["dev", "etc", "proc", "sys"] {
+        let d = tmp_dir.join(dir);
+        if d.is_dir() {
+            let _ = fs_err::remove_dir(&d);
+        }
+    }
+
+    // Verify metadata.json exists
+    if !tmp_dir.join("metadata.json").exists() {
+        let _ = fs_err::remove_dir_all(&tmp_dir);
+        bail!("pulled image does not contain metadata.json - not a valid dstack guest image");
+    }
+
+    // Rename to final location
+    fs_err::rename(&tmp_dir, &output_dir).with_context(|| {
+        format!(
+            "failed to rename {} to {}",
+            tmp_dir.display(),
+            output_dir.display()
+        )
+    })?;
+
+    info!("image extracted to {}", output_dir.display());
+    Ok(())
+}
+
+async fn extract_container(container_id: &str, dst: &Path) -> Result<()> {
+    // Use shell pipe: docker export <id> | tar x -C <dst>
+    let output = Command::new("sh")
+        .args([
+            "-c",
+            &format!(
+                "docker export {} | tar x -C {}",
+                shell_escape(container_id),
+                shell_escape(&dst.display().to_string()),
+            ),
+        ])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .output()
+        .await
+        .context("failed to run docker export | tar")?;
+
+    if !output.status.success() {
+        bail!(
+            "docker export | tar failed: {}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+    }
+
+    Ok(())
+}
+
+fn shell_escape(s: &str) -> String {
+    format!("'{}'", s.replace('\'', "'\\''"))
+}
+
+/// Determine output directory name. Read metadata.json from the image to get version,
+/// then construct the directory name as `dstack-{version}` or `dstack-{variant}-{version}`.
+async fn determine_output_dir(
+    _image_ref: &str,
+    tag: &str,
+    image_path: &Path,
+) -> Result<std::path::PathBuf> {
+    // Use tag as directory name, prefixed with "dstack-" if not already
+    let dir_name = if tag.starts_with("dstack-") {
+        tag.to_string()
+    } else {
+        format!("dstack-{tag}")
+    };
+    Ok(image_path.join(dir_name))
+}
+
+/// Parse "registry.example.com/repo/name" into ("registry.example.com", "repo/name").
+fn parse_image_ref(image_ref: &str) -> Result<(String, String)> {
+    let trimmed = image_ref
+        .trim_start_matches("https://")
+        .trim_start_matches("http://");
+
+    let first_slash = trimmed
+        .find('/')
+        .context("invalid image reference: no repository path")?;
+
+    let registry = &trimmed[..first_slash];
+    let repo = &trimmed[first_slash + 1..];
+
+    if repo.is_empty() {
+        bail!("invalid image reference: empty repository");
+    }
+
+    Ok((registry.to_string(), repo.to_string()))
+}
+
+#[derive(Deserialize)]
+struct TagList {
+    tags: Option<Vec<String>>,
+}
+
+#[derive(Deserialize)]
+struct TokenResponse {
+    token: String,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_image_ref() {
+        let (reg, repo) = parse_image_ref("cr.kvin.wang/dstack/guest-image").unwrap();
+        assert_eq!(reg, "cr.kvin.wang");
+        assert_eq!(repo, "dstack/guest-image");
+    }
+
+    #[test]
+    fn test_parse_image_ref_with_scheme() {
+        let (reg, repo) = parse_image_ref("https://ghcr.io/dstack-tee/guest-image").unwrap();
+        assert_eq!(reg, "ghcr.io");
+        assert_eq!(repo, "dstack-tee/guest-image");
+    }
+}

vmm/src/config.rs

@@ -322,6 +322,10 @@ pub struct Config {
     #[serde(default)]
     pub node_name: String,
 
+    /// OCI image registry for guest images (e.g., "cr.kvin.wang/dstack/guest-image")
+    #[serde(default)]
+    pub image_registry: String,
+
     /// The buffer size in VMM process for guest events
     pub event_buffer_size: usize,