archy/core/archipelago/src/container/registry.rs

//! Dynamic container registry configuration.
//!
//! Manages a list of container registries that the node uses to pull app images.
//! Registries are tried in order — if the first fails, the next is attempted.
//! Configuration is persisted to disk and editable via RPC.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;
use tokio::fs;
use tracing::{debug, info};

const REGISTRY_FILE: &str = "config/registries.json";

/// A single container registry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Registry {
    /// Registry URL (e.g., "git.tx1138.com/lfg2025" or "23.182.128.160:3000/lfg2025").
    pub url: String,
    /// Human-readable name.
    pub name: String,
    /// Whether TLS verification is required (false for HTTP registries).
    pub tls_verify: bool,
    /// Whether this registry is enabled.
    #[serde(default = "default_true")]
    pub enabled: bool,
    /// Priority (lower = tried first).
    #[serde(default)]
    pub priority: u32,
}

fn default_true() -> bool {
    true
}

/// Registry configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegistryConfig {
    pub registries: Vec<Registry>,
}

impl Default for RegistryConfig {
    fn default() -> Self {
        Self {
            registries: vec![
                Registry {
                    url: "23.182.128.160:3000/lfg2025".to_string(),
                    name: "Server 1 (VPS)".to_string(),
                    tls_verify: false,
                    enabled: true,
                    priority: 0,
                },
                Registry {
                    url: "git.tx1138.com/lfg2025".to_string(),
                    name: "Server 2 (tx1138)".to_string(),
                    tls_verify: true,
                    enabled: true,
                    priority: 10,
                },
            ],
        }
    }
}

impl RegistryConfig {
    /// Get enabled registries sorted by priority.
    pub fn active_registries(&self) -> Vec<&Registry> {
        let mut regs: Vec<&Registry> = self.registries.iter().filter(|r| r.enabled).collect();
        regs.sort_by_key(|r| r.priority);
        regs
    }

    /// Rewrite an image reference to use a specific registry.
    /// E.g., "git.tx1138.com/lfg2025/bitcoin-knots:latest" with registry "23.182.128.160:3000/lfg2025"
    /// becomes "23.182.128.160:3000/lfg2025/bitcoin-knots:latest".
    pub fn rewrite_image(&self, image: &str, registry: &Registry) -> String {
        // Extract the image name (last component after the org/namespace)
        // Handles: "registry/org/image:tag" -> "image:tag"
        let image_name = extract_image_name(image);
        format!("{}/{}", registry.url, image_name)
    }

    /// Generate fallback image URLs to try (excludes the original since it already failed).
    pub fn image_candidates(&self, image: &str) -> Vec<(String, bool)> {
        let mut candidates = Vec::new();

        // Rewrite for each active registry (skip if identical to original)
        for reg in self.active_registries() {
            let rewritten = self.rewrite_image(image, reg);
            if rewritten != image {
                candidates.push((rewritten, reg.tls_verify));
            }
        }

        candidates
    }

    /// Rewrite an image to use the highest-priority enabled registry, so
    /// the FIRST pull attempt honors the operator's primary choice instead
    /// of blindly using whatever registry the image URL was hardcoded to.
    /// Returns (rewritten_url, tls_verify) — or the original URL + default
    /// tls_verify=true if there's no primary (no enabled registries).
    pub fn rewrite_for_primary(&self, image: &str) -> (String, bool) {
        match self.active_registries().first() {
            Some(primary) => (self.rewrite_image(image, primary), primary.tls_verify),
            None => (image.to_string(), true),
        }
    }
}

/// Load the registry config and rewrite an image to use the primary
/// registry's URL. Convenience wrapper for callers that don't already
/// have a `RegistryConfig` in hand.
pub async fn primary_image_url(data_dir: &Path, image: &str) -> (String, bool) {
    match load_registries(data_dir).await {
        Ok(config) => config.rewrite_for_primary(image),
        Err(_) => (image.to_string(), true),
    }
}

/// Extract the image name from a full image reference.
/// "git.tx1138.com/lfg2025/bitcoin-knots:latest" -> "bitcoin-knots:latest"
/// "docker.io/gitea/gitea:1.23" -> "gitea:1.23"
fn extract_image_name(image: &str) -> &str {
    // Split by '/' and take the last segment (image:tag)
    image.rsplit('/').next().unwrap_or(image)
}

/// Load registry config from disk.
pub async fn load_registries(data_dir: &Path) -> Result<RegistryConfig> {
    let path = data_dir.join(REGISTRY_FILE);
    if !path.exists() {
        return Ok(RegistryConfig::default());
    }
    let content = fs::read_to_string(&path)
        .await
        .context("Failed to read registry config")?;
    let config: RegistryConfig =
        serde_json::from_str(&content).unwrap_or_else(|_| RegistryConfig::default());
    Ok(config)
}

/// Save registry config to disk.
pub async fn save_registries(data_dir: &Path, config: &RegistryConfig) -> Result<()> {
    let dir = data_dir.join("config");
    fs::create_dir_all(&dir)
        .await
        .context("Failed to create config dir")?;
    let path = data_dir.join(REGISTRY_FILE);
    let content =
        serde_json::to_string_pretty(config).context("Failed to serialize registry config")?;
    fs::write(&path, content)
        .await
        .context("Failed to write registry config")?;
    Ok(())
}

/// Try pulling an image from configured registries in priority order.
/// If `already_tried` is Some, that URL is skipped (avoids retrying the
/// primary when the caller already attempted it with progress streaming).
/// Returns the image reference that succeeded.
pub async fn pull_from_registries_with_skip(
    data_dir: &Path,
    image: &str,
    tmpdir: &str,
    already_tried: Option<&str>,
) -> Result<String> {
    let config = load_registries(data_dir).await?;
    let mut candidates = config.image_candidates(image);
    if let Some(skip) = already_tried {
        candidates.retain(|(url, _)| url != skip);
    }

    for (candidate, tls_verify) in &candidates {
        debug!("Trying registry: {}", candidate);

        let mut args = vec!["pull".to_string(), candidate.clone()];
        if !tls_verify {
            args.push("--tls-verify=false".to_string());
        }

        let mut child = tokio::process::Command::new("podman")
            .args(&args)
            .env("TMPDIR", tmpdir)
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .spawn()
            .ok();

        let status = if let Some(ref mut c) = child {
            match tokio::time::timeout(std::time::Duration::from_secs(120), c.wait()).await {
                Ok(Ok(s)) => Some(s.success()),
                _ => {
                    let _ = c.kill().await;
                    let _ = c.wait().await;
                    debug!("Fallback pull timed out: {}", candidate);
                    None
                }
            }
        } else {
            None
        };

        if status == Some(true) {
            // If we pulled from a non-original registry, tag it with the original name
            if candidate != image {
                let _ = tokio::process::Command::new("podman")
                    .args(["tag", candidate, image])
                    .status()
                    .await;
                info!(
                    "Pulled {} from fallback registry, tagged as {}",
                    candidate, image
                );
            } else {
                info!("Pulled {} from primary registry", image);
            }
            return Ok(candidate.clone());
        }

        debug!("Failed to pull from {}", candidate);
    }

    Err(anyhow::anyhow!(
        "Failed to pull {} from all {} configured registries",
        image,
        candidates.len()
    ))
}


#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn test_extract_image_name() {
        assert_eq!(
            extract_image_name("git.tx1138.com/lfg2025/bitcoin-knots:latest"),
            "bitcoin-knots:latest"
        );
        assert_eq!(
            extract_image_name("docker.io/gitea/gitea:1.23"),
            "gitea:1.23"
        );
        assert_eq!(extract_image_name("localhost/myimage:v1"), "myimage:v1");
    }

    #[test]
    fn test_rewrite_image() {
        let config = RegistryConfig::default();
        // Default primary is now VPS (index 0). A tx1138-hardcoded image
        // rewrites to VPS when asked for the primary mirror.
        let primary = &config.registries[0];
        assert_eq!(
            config.rewrite_image("git.tx1138.com/lfg2025/bitcoin-knots:latest", primary),
            "23.182.128.160:3000/lfg2025/bitcoin-knots:latest"
        );
    }

    #[test]
    fn test_image_candidates() {
        let config = RegistryConfig::default();
        let candidates = config.image_candidates("git.tx1138.com/lfg2025/lnd:v0.18.4-beta");
        // Defaults: VPS (primary) + tx1138. tx1138 is filtered out because
        // it's identical to the original image URL, leaving one candidate.
        assert_eq!(candidates.len(), 1);
        // Primary-first — VPS rewrite leads the candidate list.
        assert_eq!(candidates[0].0, "23.182.128.160:3000/lfg2025/lnd:v0.18.4-beta");
    }

    #[test]
    fn test_rewrite_for_primary_uses_top_priority() {
        let config = RegistryConfig::default();
        let (url, tls) =
            config.rewrite_for_primary("git.tx1138.com/lfg2025/lnd:v0.18.4-beta");
        assert_eq!(url, "23.182.128.160:3000/lfg2025/lnd:v0.18.4-beta");
        assert!(!tls, "VPS primary is HTTP — tls_verify should be false");
    }

    #[test]
    fn test_active_registries_sorted() {
        let config = RegistryConfig::default();
        let active = config.active_registries();
        assert_eq!(active.len(), 2);
        assert!(active[0].priority <= active[1].priority);
    }

    #[tokio::test]
    async fn test_load_default() {
        let tmp = TempDir::new().unwrap();
        let config = load_registries(tmp.path()).await.unwrap();
        assert_eq!(config.registries.len(), 2);
    }

    #[tokio::test]
    async fn test_save_load_roundtrip() {
        let tmp = TempDir::new().unwrap();
        let mut config = RegistryConfig::default();
        config.registries.push(Registry {
            url: "myregistry.com/apps".into(),
            name: "Custom".into(),
            tls_verify: true,
            enabled: true,
            priority: 5,
        });
        save_registries(tmp.path(), &config).await.unwrap();
        let loaded = load_registries(tmp.path()).await.unwrap();
        assert_eq!(loaded.registries.len(), 3);
    }
}