zlayer_core/auth/

resolver.rs

//! Authentication resolver for OCI registries
//!
//! This module provides flexible authentication resolution supporting multiple sources
//! and per-registry configuration.
//!
//! The wire DTOs (`AuthSource`, `AuthConfig`, `RegistryAuthConfig`) now live in
//! `zlayer-types::auth` and are re-exported here so existing
//! `zlayer_core::auth::AuthSource` (etc.) import paths keep working.

pub use zlayer_types::auth::{AuthConfig, AuthSource, RegistryAuthConfig};

use super::DockerConfigAuth;
use std::collections::HashMap;
use std::path::PathBuf;

/// Authentication resolver that converts `AuthConfig` to `oci_client` `RegistryAuth`
pub struct AuthResolver {
    config: AuthConfig,
    docker_config: Option<DockerConfigAuth>,
    registry_map: HashMap<String, AuthSource>,
}

impl AuthResolver {
    /// Create a new authentication resolver
    #[must_use]
    pub fn new(config: AuthConfig) -> Self {
        // Build a map for fast registry lookups
        let registry_map: HashMap<String, AuthSource> = config
            .registries
            .iter()
            .map(|r| (r.registry.clone(), r.source.clone()))
            .collect();

        // Load Docker config if any source uses DockerConfig
        let needs_docker_config = config.default == AuthSource::DockerConfig
            || registry_map
                .values()
                .any(|s| matches!(s, AuthSource::DockerConfig));

        let docker_config = if needs_docker_config {
            Self::load_docker_config(config.docker_config_path.as_ref())
        } else {
            None
        };

        Self {
            config,
            docker_config,
            registry_map,
        }
    }

    /// Resolve authentication for an image reference
    ///
    /// Extracts the registry from the image reference and returns the appropriate
    /// `oci_client::secrets::RegistryAuth`.
    #[must_use]
    pub fn resolve(&self, image: &str) -> oci_client::secrets::RegistryAuth {
        let registry = Self::extract_registry(image);
        let source = self
            .registry_map
            .get(&registry)
            .unwrap_or(&self.config.default);

        self.resolve_source(source, &registry)
    }

    /// Return the `AuthSource` that would be used for the given registry hostname.
    ///
    /// Looks up the per-registry map first, falling back to the default source.
    #[must_use]
    pub fn source_for_registry(&self, registry: &str) -> &AuthSource {
        self.registry_map
            .get(registry)
            .unwrap_or(&self.config.default)
    }

    /// Resolve a specific `AuthSource` to `RegistryAuth`.
    ///
    /// This is the synchronous resolution path. `AuthSource::SecretStore`
    /// cannot be resolved synchronously and returns `Anonymous` with a
    /// warning log.
    pub fn resolve_source(
        &self,
        source: &AuthSource,
        registry: &str,
    ) -> oci_client::secrets::RegistryAuth {
        match source {
            AuthSource::Anonymous => oci_client::secrets::RegistryAuth::Anonymous,

            AuthSource::Basic { username, password } => {
                oci_client::secrets::RegistryAuth::Basic(username.clone(), password.clone())
            }

            AuthSource::DockerConfig => {
                if let Some(ref docker_config) = self.docker_config {
                    if let Some((username, password)) = docker_config.get_credentials(registry) {
                        return oci_client::secrets::RegistryAuth::Basic(username, password);
                    }
                }
                // Fallback to anonymous if no credentials found
                oci_client::secrets::RegistryAuth::Anonymous
            }

            AuthSource::EnvVar {
                username_var,
                password_var,
            } => {
                let username = std::env::var(username_var).unwrap_or_default();
                let password = std::env::var(password_var).unwrap_or_default();

                if !username.is_empty() && !password.is_empty() {
                    oci_client::secrets::RegistryAuth::Basic(username, password)
                } else {
                    oci_client::secrets::RegistryAuth::Anonymous
                }
            }

            AuthSource::SecretStore { .. } => {
                tracing::warn!(
                    "SecretStore auth source requires async resolver; returning Anonymous"
                );
                oci_client::secrets::RegistryAuth::Anonymous
            }
        }
    }

    /// Normalize a registry hostname for case-insensitive, scheme/slash
    /// insensitive matching.
    ///
    /// Strips a leading `http://` / `https://`, trims trailing `/`, and
    /// lowercases. This lets a credential stored under `https://ghcr.io/`
    /// (as a user might type it into `zlayer login`) match an image whose
    /// extracted registry is `ghcr.io`.
    ///
    /// Examples:
    /// - `https://ghcr.io/` -> `ghcr.io`
    /// - `GHCR.IO` -> `ghcr.io`
    /// - `localhost:5000` -> `localhost:5000`
    #[must_use]
    pub fn normalize_registry(registry: &str) -> String {
        let s = registry.trim();
        let s = s
            .strip_prefix("https://")
            .or_else(|| s.strip_prefix("http://"))
            .unwrap_or(s);
        s.trim_end_matches('/').to_ascii_lowercase()
    }

    /// Extract registry hostname from image reference
    ///
    /// Examples:
    /// - "ubuntu:latest" -> "docker.io"
    /// - "ghcr.io/owner/repo:tag" -> "ghcr.io"
    /// - "localhost:5000/image" -> "localhost:5000"
    #[must_use]
    pub fn extract_registry(image: &str) -> String {
        // Remove digest if present
        let image_without_digest = image.split('@').next().unwrap_or(image);

        // Split by '/'
        let parts: Vec<&str> = image_without_digest.split('/').collect();

        // If there's no '/', it's just an image name, assume Docker Hub
        if parts.len() == 1 {
            return "docker.io".to_string();
        }

        // Check if first part looks like a hostname (contains '.' or ':' or is 'localhost')
        let first_part = parts[0];
        if first_part.contains('.') || first_part.contains(':') || first_part == "localhost" {
            first_part.to_string()
        } else {
            // No explicit registry (e.g., "library/ubuntu"), assume Docker Hub
            "docker.io".to_string()
        }
    }

    /// Load Docker config from path or default location
    fn load_docker_config(path: Option<&PathBuf>) -> Option<DockerConfigAuth> {
        let config = if let Some(path) = path {
            DockerConfigAuth::load_from_path(path).ok()
        } else {
            DockerConfigAuth::load().ok()
        };

        if config.is_none() {
            tracing::debug!("Failed to load Docker config, using anonymous auth as fallback");
        }

        config
    }
}

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

    #[test]
    fn test_extract_registry() {
        assert_eq!(AuthResolver::extract_registry("ubuntu"), "docker.io");
        assert_eq!(AuthResolver::extract_registry("ubuntu:latest"), "docker.io");
        assert_eq!(
            AuthResolver::extract_registry("library/ubuntu"),
            "docker.io"
        );
        assert_eq!(
            AuthResolver::extract_registry("ghcr.io/owner/repo"),
            "ghcr.io"
        );
        assert_eq!(
            AuthResolver::extract_registry("ghcr.io/owner/repo:tag"),
            "ghcr.io"
        );
        assert_eq!(
            AuthResolver::extract_registry("localhost:5000/image"),
            "localhost:5000"
        );
        assert_eq!(
            AuthResolver::extract_registry("myregistry.com/path/to/image:v1.0"),
            "myregistry.com"
        );
    }

    #[test]
    fn test_normalize_registry() {
        assert_eq!(AuthResolver::normalize_registry("ghcr.io"), "ghcr.io");
        assert_eq!(AuthResolver::normalize_registry("GHCR.IO"), "ghcr.io");
        assert_eq!(
            AuthResolver::normalize_registry("https://ghcr.io"),
            "ghcr.io"
        );
        assert_eq!(
            AuthResolver::normalize_registry("https://ghcr.io/"),
            "ghcr.io"
        );
        assert_eq!(
            AuthResolver::normalize_registry("http://localhost:5000/"),
            "localhost:5000"
        );
        assert_eq!(AuthResolver::normalize_registry("  ghcr.io  "), "ghcr.io");
    }

    #[test]
    fn test_anonymous_auth() {
        let config = AuthConfig {
            default: AuthSource::Anonymous,
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);
        let auth = resolver.resolve("ubuntu:latest");

        assert!(matches!(auth, oci_client::secrets::RegistryAuth::Anonymous));
    }

    #[test]
    fn test_basic_auth() {
        let config = AuthConfig {
            default: AuthSource::Basic {
                username: "user".to_string(),
                password: "pass".to_string(),
            },
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);
        let auth = resolver.resolve("ubuntu:latest");

        match auth {
            oci_client::secrets::RegistryAuth::Basic(username, password) => {
                assert_eq!(username, "user");
                assert_eq!(password, "pass");
            }
            _ => panic!("Expected Basic auth"),
        }
    }

    #[test]
    fn test_per_registry_auth() {
        let config = AuthConfig {
            registries: vec![RegistryAuthConfig {
                registry: "ghcr.io".to_string(),
                source: AuthSource::Basic {
                    username: "ghcr_user".to_string(),
                    password: "ghcr_pass".to_string(),
                },
            }],
            default: AuthSource::Anonymous,
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);

        // Should use specific auth for ghcr.io
        let auth = resolver.resolve("ghcr.io/owner/repo:tag");
        match auth {
            oci_client::secrets::RegistryAuth::Basic(username, password) => {
                assert_eq!(username, "ghcr_user");
                assert_eq!(password, "ghcr_pass");
            }
            _ => panic!("Expected Basic auth for ghcr.io"),
        }

        // Should use default (anonymous) for docker.io
        let auth = resolver.resolve("ubuntu:latest");
        assert!(matches!(auth, oci_client::secrets::RegistryAuth::Anonymous));
    }

    #[test]
    fn test_env_var_auth() {
        std::env::set_var("TEST_USERNAME", "env_user");
        std::env::set_var("TEST_PASSWORD", "env_pass");

        let config = AuthConfig {
            default: AuthSource::EnvVar {
                username_var: "TEST_USERNAME".to_string(),
                password_var: "TEST_PASSWORD".to_string(),
            },
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);
        let auth = resolver.resolve("ubuntu:latest");

        match auth {
            oci_client::secrets::RegistryAuth::Basic(username, password) => {
                assert_eq!(username, "env_user");
                assert_eq!(password, "env_pass");
            }
            _ => panic!("Expected Basic auth from env vars"),
        }

        std::env::remove_var("TEST_USERNAME");
        std::env::remove_var("TEST_PASSWORD");
    }

    #[test]
    fn test_env_var_auth_fallback() {
        // Test that missing env vars fall back to anonymous
        let config = AuthConfig {
            default: AuthSource::EnvVar {
                username_var: "NONEXISTENT_USER".to_string(),
                password_var: "NONEXISTENT_PASS".to_string(),
            },
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);
        let auth = resolver.resolve("ubuntu:latest");

        assert!(matches!(auth, oci_client::secrets::RegistryAuth::Anonymous));
    }

    #[test]
    fn test_secret_store_sync_fallback_returns_anonymous() {
        let config = AuthConfig {
            registries: vec![RegistryAuthConfig {
                registry: "private.registry.io".to_string(),
                source: AuthSource::SecretStore {
                    credential_id: "cred-uuid-123".to_string(),
                },
            }],
            default: AuthSource::Anonymous,
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);

        // The sync path cannot resolve SecretStore and must return Anonymous
        let auth = resolver.resolve("private.registry.io/image:latest");
        assert!(matches!(auth, oci_client::secrets::RegistryAuth::Anonymous));

        // The default source should still work normally
        let auth = resolver.resolve("ubuntu:latest");
        assert!(matches!(auth, oci_client::secrets::RegistryAuth::Anonymous));
    }

    #[test]
    fn test_source_for_registry_returns_correct_source() {
        let config = AuthConfig {
            registries: vec![RegistryAuthConfig {
                registry: "ghcr.io".to_string(),
                source: AuthSource::Basic {
                    username: "user".to_string(),
                    password: "pass".to_string(),
                },
            }],
            default: AuthSource::Anonymous,
            ..Default::default()
        };

        let resolver = AuthResolver::new(config);

        // Known registry returns its configured source
        let source = resolver.source_for_registry("ghcr.io");
        assert!(matches!(source, AuthSource::Basic { .. }));

        // Unknown registry returns the default
        let source = resolver.source_for_registry("docker.io");
        assert!(matches!(source, AuthSource::Anonymous));
    }

    #[test]
    fn test_secret_store_serde_roundtrip() {
        let source = AuthSource::SecretStore {
            credential_id: "abc-123".to_string(),
        };
        let json = serde_json::to_string(&source).unwrap();
        let parsed: AuthSource = serde_json::from_str(&json).unwrap();
        assert_eq!(source, parsed);
    }
}