photoncloud-monorepo/deployer/crates/deployer-server/src/cloud_init.rs

use axum::{
    extract::{Path, State},
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
};
use deployer_types::NodeConfig;
use std::sync::Arc;

use crate::{
    auth::require_bootstrap_auth, phone_home::lookup_node_config, state::AppState,
    validation::validate_identifier,
};

/// GET /api/v1/cloud-init/:machine_id/meta-data
pub async fn meta_data(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    Path(machine_id): Path<String>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    require_bootstrap_auth(&state, &headers)?;
    validate_identifier(&machine_id, "machine_id")?;

    let Some((node_id, config)) = lookup_node_config(&state, &machine_id).await else {
        return Err((
            StatusCode::NOT_FOUND,
            "machine-id not registered".to_string(),
        ));
    };

    let body = format!(
        "instance-id: {}\nlocal-hostname: {}\n",
        node_id, config.hostname
    );
    Ok(([(axum::http::header::CONTENT_TYPE, "text/plain")], body))
}

/// GET /api/v1/cloud-init/:machine_id/user-data
pub async fn user_data(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
    Path(machine_id): Path<String>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    require_bootstrap_auth(&state, &headers)?;
    validate_identifier(&machine_id, "machine_id")?;

    let Some((node_id, config)) = lookup_node_config(&state, &machine_id).await else {
        return Err((
            StatusCode::NOT_FOUND,
            "machine-id not registered".to_string(),
        ));
    };

    let body = render_user_data(&node_id, &config)
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
    Ok((
        [(axum::http::header::CONTENT_TYPE, "text/cloud-config")],
        body,
    ))
}

fn render_yaml_list(items: &[String], indent: usize) -> String {
    if items.is_empty() {
        return "[]".to_string();
    }

    let prefix = " ".repeat(indent);
    items
        .iter()
        .map(|item| format!("{prefix}- {:?}", item))
        .collect::<Vec<_>>()
        .join("\n")
}

fn indent_multiline(input: &str, indent: usize) -> String {
    let prefix = " ".repeat(indent);
    input
        .lines()
        .map(|line| format!("{prefix}{line}"))
        .collect::<Vec<_>>()
        .join("\n")
}

fn render_user_data(node_id: &str, config: &NodeConfig) -> anyhow::Result<String> {
    let node_config_json = serde_json::to_string_pretty(config)?;
    let ssh_keys = render_yaml_list(&config.ssh_authorized_keys, 2);

    Ok(format!(
        r#"#cloud-config
hostname: {hostname}
fqdn: {hostname}
manage_etc_hosts: true
ssh_authorized_keys:
{ssh_keys}
write_files:
  - path: /etc/plasmacloud/node-id
    permissions: "0644"
    content: |
{node_id_block}
  - path: /etc/plasmacloud/node-config.json
    permissions: "0644"
    content: |
{node_config_block}
"#,
        hostname = config.hostname,
        ssh_keys = ssh_keys,
        node_id_block = indent_multiline(node_id, 6),
        node_config_block = indent_multiline(&node_config_json, 6),
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;
    use crate::state::AppState;
    use axum::body::Body;
    use axum::http::Request;
    use deployer_types::InstallPlan;
    use tower::ServiceExt;

    fn test_config() -> NodeConfig {
        NodeConfig {
            hostname: "node01".to_string(),
            role: "worker".to_string(),
            ip: "10.0.0.11".to_string(),
            services: vec!["prismnet".to_string()],
            ssh_authorized_keys: vec!["ssh-ed25519 AAAATEST test".to_string()],
            labels: std::collections::HashMap::from([("tier".to_string(), "general".to_string())]),
            pool: Some("general".to_string()),
            node_class: Some("worker".to_string()),
            failure_domain: Some("rack-a".to_string()),
            nix_profile: Some("profiles/worker".to_string()),
            install_plan: Some(InstallPlan {
                nixos_configuration: Some("worker-golden".to_string()),
                disko_config_path: Some("profiles/worker/disko.nix".to_string()),
                target_disk: Some("/dev/vda".to_string()),
                target_disk_by_id: None,
            }),
        }
    }

    #[test]
    fn test_render_user_data_contains_node_config() {
        let rendered = render_user_data("node01", &test_config()).unwrap();
        assert!(rendered.contains("#cloud-config"));
        assert!(rendered.contains("hostname: node01"));
        assert!(rendered.contains("/etc/plasmacloud/node-config.json"));
        assert!(rendered.contains("\"nix_profile\": \"profiles/worker\""));
        assert!(rendered.contains("\"nixos_configuration\": \"worker-golden\""));
    }

    #[tokio::test]
    async fn test_cloud_init_routes() {
        let mut config = Config::default();
        config.bootstrap_token = Some("test-token".to_string());
        let state = Arc::new(AppState::with_config(config));
        state.machine_configs.write().await.insert(
            "machine-1".to_string(),
            ("node01".to_string(), test_config()),
        );
        let app = crate::build_router(state);

        let response = app
            .clone()
            .oneshot(
                Request::builder()
                    .uri("/api/v1/cloud-init/machine-1/meta-data")
                    .header("x-deployer-token", "test-token")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/api/v1/cloud-init/machine-1/user-data")
                    .header("x-deployer-token", "test-token")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }
}