mockforge_registry_server/deployment/

flyio.rs

//! Fly.io API integration for deploying mock services

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Fly.io API client for managing deployments
pub struct FlyioClient {
    api_token: String,
    base_url: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioApp {
    pub id: String,
    #[serde(default)]
    pub name: Option<String>,
    #[serde(default)]
    pub hostname: Option<String>,
    #[serde(default)]
    pub organization: Option<FlyioOrganization>,
    #[serde(default)]
    pub status: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioOrganization {
    pub id: String,
    #[serde(default)]
    pub name: Option<String>,
    #[serde(default)]
    pub slug: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioMachine {
    pub id: String,
    pub name: String,
    pub state: String,
    pub region: String,
    pub image_ref: Option<FlyioImageRef>,
    pub config: FlyioMachineConfig,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioImageRef {
    pub registry: String,
    pub repository: String,
    pub tag: String,
    pub digest: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioMachineConfig {
    pub image: String,
    pub env: HashMap<String, String>,
    pub services: Vec<FlyioService>,
    pub checks: Option<HashMap<String, FlyioCheck>>,
    /// VM resource sizing. `None` means accept Fly's API default
    /// (currently `shared-cpu-1x:256MB`); existing deployments and
    /// non-plugin-enabled machines stay on that default. Cloud
    /// plugins requires explicit sizing — see
    /// `docs/plugins/security/cloud-runtime-sidecar-spike.md` for
    /// the measured-on-Fly numbers and tier table.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub guest: Option<FlyioGuest>,
}

/// VM sizing block sent to Fly Machines API as `config.guest`.
/// Mirrors the API's expected JSON shape:
/// <https://fly.io/docs/machines/api/machines-resource/#machine-config-object-properties>
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlyioGuest {
    /// `"shared"` or `"performance"`.
    pub cpu_kind: String,
    pub cpus: u32,
    pub memory_mb: u32,
}

impl FlyioGuest {
    /// `shared-cpu-1x:256MB` — Fly's default for the registry today.
    /// Used when no plugins are attached so existing hosted-mocks
    /// keep their current footprint.
    pub fn shared_256() -> Self {
        Self {
            cpu_kind: "shared".into(),
            cpus: 1,
            memory_mb: 256,
        }
    }

    /// `shared-cpu-1x:512MB` — Pro tier with cloud plugins. Real-Fly
    /// measurement on the spike showed mockforge + sidecar idle at
    /// ~166 MB / 256 MB (65%), with only 71 MB headroom; bumping to
    /// 512 MB gives 316 MB headroom — comfortable for ≤5 plugins
    /// at typical sizes.
    pub fn shared_512() -> Self {
        Self {
            cpu_kind: "shared".into(),
            cpus: 1,
            memory_mb: 512,
        }
    }

    /// `shared-cpu-1x:1024MB` — Team tier with cloud plugins
    /// (≤25 plugins per mock).
    pub fn shared_1024() -> Self {
        Self {
            cpu_kind: "shared".into(),
            cpus: 1,
            memory_mb: 1024,
        }
    }

    /// `shared-cpu-2x:2048MB` — Enterprise tier with metered
    /// pass-through. Two cores so heavy plugin workloads don't
    /// starve mockforge's request path.
    pub fn shared_2x_2048() -> Self {
        Self {
            cpu_kind: "shared".into(),
            cpus: 2,
            memory_mb: 2048,
        }
    }

    /// Pick the right Fly machine size for a hosted-mock deployment
    /// based on org plan + whether cloud plugins are attached.
    ///
    /// Without plugins, every tier stays on the existing 256 MB
    /// default — no behavior change for legacy deployments.
    ///
    /// With plugins, the floor bumps according to the
    /// real-microVM measurements in
    /// `docs/plugins/security/cloud-runtime-sidecar-spike.md`:
    /// idle on 256 MB sits at ~166 MB / 65% utilization with only
    /// ~71 MB headroom. That's not enough for the plugin runtime to
    /// grow into without OOM-killing mockforge. 512 MB gives 316 MB
    /// headroom which fits ≤5 plugins comfortably; Team's larger
    /// plugin count needs 1024 MB.
    pub fn for_hosted_mock(plan: &str, plugins_enabled: bool) -> Self {
        if !plugins_enabled {
            return Self::shared_256();
        }
        match plan.to_lowercase().as_str() {
            "free" => Self::shared_256(), // Plugins not available on Free.
            "pro" => Self::shared_512(),
            "team" => Self::shared_1024(),
            // Unknown / future plans (e.g. "enterprise") get the
            // largest currently-available shape. The handler should
            // refuse plugin attachment for unknown plans before
            // ever reaching this; defaulting big is fail-safe.
            _ => Self::shared_2x_2048(),
        }
    }
}

/// Registry authentication for pulling private Docker images
#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioRegistryAuth {
    pub server: String,
    pub username: String,
    pub password: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioService {
    pub protocol: String,
    pub internal_port: u16,
    pub ports: Vec<FlyioPort>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioPort {
    pub port: u16,
    pub handlers: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FlyioCheck {
    #[serde(rename = "type")]
    pub check_type: String,
    pub port: u16,
    pub grace_period: String,
    pub interval: String,
    pub method: String,
    pub timeout: String,
    pub tls_skip_verify: bool,
    pub path: Option<String>,
}

impl FlyioClient {
    pub fn new(api_token: String) -> Self {
        Self {
            api_token,
            base_url: "https://api.machines.dev".to_string(),
        }
    }

    /// Get the API token (used for Fly.io registry auth)
    pub fn api_token(&self) -> &str {
        &self.api_token
    }

    /// Create a new Fly.io app
    pub async fn create_app(&self, app_name: &str, org_slug: &str) -> Result<FlyioApp> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps", self.base_url);

        let payload = serde_json::json!({
            "app_name": app_name,
            "org_slug": org_slug,
        });

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .header("Content-Type", "application/json")
            .json(&payload)
            .send()
            .await
            .context("Failed to create Fly.io app")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());

            // Handle "Name has already been taken" (422) — the app exists, fetch it instead
            if status.as_u16() == 422 && error_text.contains("already been taken") {
                tracing::info!("Fly.io app '{}' already exists, fetching existing app", app_name);
                return self.get_app(app_name).await;
            }

            anyhow::bail!("Failed to create Fly.io app: {} - {}", status, error_text);
        }

        let app: FlyioApp = response.json().await.context("Failed to parse Fly.io app response")?;

        Ok(app)
    }

    /// Create a machine (instance) for the app
    pub async fn create_machine(
        &self,
        app_name: &str,
        config: FlyioMachineConfig,
        region: &str,
        registry_auth: Option<FlyioRegistryAuth>,
    ) -> Result<FlyioMachine> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines", self.base_url, app_name);

        let mut payload = serde_json::json!({
            "config": config,
            "region": region,
        });
        if let Some(auth) = registry_auth {
            payload["config"]["image_registry_auth"] =
                serde_json::to_value(auth).context("Failed to serialize registry auth")?;
        }

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .header("Content-Type", "application/json")
            .json(&payload)
            .send()
            .await
            .context("Failed to create Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to create Fly.io machine: {} - {}", status, error_text);
        }

        let machine: FlyioMachine =
            response.json().await.context("Failed to parse Fly.io machine response")?;

        Ok(machine)
    }

    /// Update a machine's configuration (image, env vars, etc.)
    pub async fn update_machine(
        &self,
        app_name: &str,
        machine_id: &str,
        config: FlyioMachineConfig,
        registry_auth: Option<FlyioRegistryAuth>,
    ) -> Result<FlyioMachine> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines/{}", self.base_url, app_name, machine_id);

        let mut payload = serde_json::json!({
            "config": config,
        });
        if let Some(auth) = registry_auth {
            payload["config"]["image_registry_auth"] =
                serde_json::to_value(auth).context("Failed to serialize registry auth")?;
        }

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .header("Content-Type", "application/json")
            .json(&payload)
            .send()
            .await
            .context("Failed to update Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to update Fly.io machine: {} - {}", status, error_text);
        }

        let machine: FlyioMachine =
            response.json().await.context("Failed to parse Fly.io machine response")?;

        Ok(machine)
    }

    /// Get machine status
    pub async fn get_machine(&self, app_name: &str, machine_id: &str) -> Result<FlyioMachine> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines/{}", self.base_url, app_name, machine_id);

        let response = client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .send()
            .await
            .context("Failed to get Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to get Fly.io machine: {} - {}", status, error_text);
        }

        let machine: FlyioMachine =
            response.json().await.context("Failed to parse Fly.io machine response")?;

        Ok(machine)
    }

    /// Stop a running machine (graceful shutdown, keeps the machine around
    /// so it can be restarted later without recreating its config).
    pub async fn stop_machine(&self, app_name: &str, machine_id: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines/{}/stop", self.base_url, app_name, machine_id);

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .header("Content-Type", "application/json")
            .json(&serde_json::json!({}))
            .send()
            .await
            .context("Failed to stop Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to stop Fly.io machine: {} - {}", status, error_text);
        }

        Ok(())
    }

    /// Start a stopped machine.
    pub async fn start_machine(&self, app_name: &str, machine_id: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines/{}/start", self.base_url, app_name, machine_id);

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .header("Content-Type", "application/json")
            .send()
            .await
            .context("Failed to start Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to start Fly.io machine: {} - {}", status, error_text);
        }

        Ok(())
    }

    /// Delete a machine
    pub async fn delete_machine(&self, app_name: &str, machine_id: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines/{}", self.base_url, app_name, machine_id);

        let response = client
            .delete(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .send()
            .await
            .context("Failed to delete Fly.io machine")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to delete Fly.io machine: {} - {}", status, error_text);
        }

        Ok(())
    }

    /// Delete a Fly.io app
    pub async fn delete_app(&self, app_name: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}", self.base_url, app_name);

        let response = client
            .delete(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .send()
            .await
            .context("Failed to delete Fly.io app")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to delete Fly.io app: {} - {}", status, error_text);
        }

        Ok(())
    }

    /// Allocate a shared IPv4 and a dedicated IPv6 address for an app
    pub async fn allocate_ips(&self, app_name: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let graphql_url = "https://api.fly.io/graphql";

        // Allocate shared IPv4
        let ipv4_query = serde_json::json!({
            "query": "mutation($input: AllocateIPAddressInput!) { allocateIpAddress(input: $input) { ipAddress { id address type } } }",
            "variables": {
                "input": {
                    "appId": app_name,
                    "type": "shared_v4"
                }
            }
        });

        let response = client
            .post(graphql_url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .json(&ipv4_query)
            .send()
            .await
            .context("Failed to allocate shared IPv4")?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Failed to allocate shared IPv4: {}", error_text);
        }

        // Allocate IPv6
        let ipv6_query = serde_json::json!({
            "query": "mutation($input: AllocateIPAddressInput!) { allocateIpAddress(input: $input) { ipAddress { id address type } } }",
            "variables": {
                "input": {
                    "appId": app_name,
                    "type": "v6"
                }
            }
        });

        let response = client
            .post(graphql_url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .json(&ipv6_query)
            .send()
            .await
            .context("Failed to allocate IPv6")?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Failed to allocate IPv6: {}", error_text);
        }

        Ok(())
    }

    /// Get app info
    pub async fn get_app(&self, app_name: &str) -> Result<FlyioApp> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}", self.base_url, app_name);

        let response = client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .send()
            .await
            .context("Failed to get Fly.io app")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to get Fly.io app: {} - {}", status, error_text);
        }

        let app: FlyioApp = response.json().await.context("Failed to parse Fly.io app response")?;

        Ok(app)
    }

    /// Add a custom domain certificate to an app
    ///
    /// This tells Fly.io to provision a Let's Encrypt TLS certificate for the
    /// given hostname and route traffic for that hostname to this app via SNI.
    pub async fn add_certificate(&self, app_name: &str, hostname: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let graphql_url = "https://api.fly.io/graphql";

        let query = serde_json::json!({
            "query": "mutation($appId: ID!, $hostname: String!) { addCertificate(appId: $appId, hostname: $hostname) { certificate { id hostname } } }",
            "variables": {
                "appId": app_name,
                "hostname": hostname
            }
        });

        let response = client
            .post(graphql_url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .json(&query)
            .send()
            .await
            .context("Failed to add certificate")?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Failed to add certificate for {}: {}", hostname, error_text);
        }

        // Check for GraphQL-level errors
        let body: serde_json::Value =
            response.json().await.context("Failed to parse certificate response")?;
        if let Some(errors) = body.get("errors") {
            // "already exists" is fine — idempotent
            let err_str = errors.to_string();
            if !err_str.contains("already exists") {
                anyhow::bail!("Failed to add certificate for {}: {}", hostname, err_str);
            }
        }

        Ok(())
    }

    /// Remove a custom domain certificate from an app
    pub async fn delete_certificate(&self, app_name: &str, hostname: &str) -> Result<()> {
        let client = reqwest::Client::new();
        let graphql_url = "https://api.fly.io/graphql";

        let query = serde_json::json!({
            "query": "mutation($appId: ID!, $hostname: String!) { deleteCertificate(appId: $appId, hostname: $hostname) { app { name } } }",
            "variables": {
                "appId": app_name,
                "hostname": hostname
            }
        });

        let response = client
            .post(graphql_url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .json(&query)
            .send()
            .await
            .context("Failed to delete certificate")?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Failed to delete certificate for {}: {}", hostname, error_text);
        }

        Ok(())
    }

    /// List machines for an app
    pub async fn list_machines(&self, app_name: &str) -> Result<Vec<FlyioMachine>> {
        let client = reqwest::Client::new();
        let url = format!("{}/v1/apps/{}/machines", self.base_url, app_name);

        let response = client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.api_token))
            .send()
            .await
            .context("Failed to list Fly.io machines")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            anyhow::bail!("Failed to list Fly.io machines: {} - {}", status, error_text);
        }

        let machines: Vec<FlyioMachine> =
            response.json().await.context("Failed to parse Fly.io machines response")?;

        Ok(machines)
    }
}

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

    #[test]
    fn for_hosted_mock_no_plugins_uses_legacy_256() {
        // Existing hosted-mocks without cloud plugins must keep
        // their current 256 MB footprint — no surprise pricing
        // bumps for legacy deployments.
        for plan in ["free", "pro", "team", "enterprise", "weird-future-plan"] {
            let g = FlyioGuest::for_hosted_mock(plan, false);
            assert_eq!(g.memory_mb, 256, "plan {} without plugins must stay at 256MB", plan);
            assert_eq!(g.cpu_kind, "shared");
            assert_eq!(g.cpus, 1);
        }
    }

    #[test]
    fn for_hosted_mock_pro_with_plugins_bumps_to_512() {
        let g = FlyioGuest::for_hosted_mock("pro", true);
        assert_eq!(g.memory_mb, 512);
        assert_eq!(g.cpus, 1);
    }

    #[test]
    fn for_hosted_mock_team_with_plugins_bumps_to_1024() {
        let g = FlyioGuest::for_hosted_mock("team", true);
        assert_eq!(g.memory_mb, 1024);
        assert_eq!(g.cpus, 1);
    }

    #[test]
    fn for_hosted_mock_free_with_plugins_stays_256() {
        // Plugins aren't available on Free — handler rejects before
        // we get here — but if we do, don't accidentally upsize
        // a Free org's machine.
        let g = FlyioGuest::for_hosted_mock("free", true);
        assert_eq!(g.memory_mb, 256);
    }

    #[test]
    fn for_hosted_mock_unknown_plan_with_plugins_fails_safe_high() {
        // Unknown plans (e.g. future "enterprise") get the largest
        // shape rather than under-provisioning.
        let g = FlyioGuest::for_hosted_mock("enterprise", true);
        assert_eq!(g.memory_mb, 2048);
        assert_eq!(g.cpus, 2);
    }

    #[test]
    fn for_hosted_mock_plan_string_is_case_insensitive() {
        let g = FlyioGuest::for_hosted_mock("PRO", true);
        assert_eq!(g.memory_mb, 512);
        let g = FlyioGuest::for_hosted_mock("Team", true);
        assert_eq!(g.memory_mb, 1024);
    }

    #[test]
    fn machine_config_serialization_omits_guest_when_none() {
        // Existing call sites that pass `guest: None` (or use the
        // legacy struct without the field via Default) must produce
        // JSON without a `guest` key — so Fly's API default kicks in.
        let cfg = FlyioMachineConfig {
            image: "img".into(),
            env: HashMap::new(),
            services: vec![],
            checks: None,
            guest: None,
        };
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(!json.contains("guest"), "guest=None must be omitted, got {}", json);
    }

    #[test]
    fn machine_config_serializes_guest_when_set() {
        let cfg = FlyioMachineConfig {
            image: "img".into(),
            env: HashMap::new(),
            services: vec![],
            checks: None,
            guest: Some(FlyioGuest::shared_512()),
        };
        let json = serde_json::to_string(&cfg).unwrap();
        assert!(json.contains("\"guest\""));
        assert!(json.contains("\"memory_mb\":512"));
        assert!(json.contains("\"cpu_kind\":\"shared\""));
    }
}