fleetforge_control_plane/

lib.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::{Context, Result};
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use chrono::{DateTime, Utc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tower_http::trace::TraceLayer;
use tracing::{info, instrument};
use uuid::Uuid;

const VALID_ROLES: &[&str] = &["admin", "viewer", "billing"];

#[derive(Debug, Clone)]
pub struct ControlPlaneConfig {
    pub addr: SocketAddr,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Tenant {
    pub id: Uuid,
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub contact: Option<String>,
    pub status: &'static str,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Project {
    pub id: Uuid,
    pub tenant_id: Uuid,
    pub name: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ApiKey {
    pub id: Uuid,
    pub project_id: Uuid,
    pub key: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TenantQuota {
    pub tenant_id: Uuid,
    #[serde(default)]
    pub daily_token_limit: Option<i64>,
    #[serde(default)]
    pub daily_cost_limit: Option<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BillingRecord {
    pub tenant_id: Uuid,
    pub recorded_at: DateTime<Utc>,
    pub tokens: i64,
    pub cost: f64,
}

#[derive(Default)]
struct ControlPlaneState {
    tenants: RwLock<HashMap<Uuid, Tenant>>,
    projects: RwLock<HashMap<Uuid, Project>>,
    keys: RwLock<HashMap<Uuid, ApiKey>>,
    quotas: RwLock<HashMap<Uuid, TenantQuota>>,
    tenant_roles: RwLock<HashMap<Uuid, HashMap<String, String>>>,
    billing: RwLock<Vec<BillingRecord>>,
}

#[derive(Debug, Deserialize)]
struct TenantSpec {
    name: String,
    #[serde(default)]
    contact: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ProjectSpec {
    name: String,
}

#[derive(Debug, Deserialize)]
struct QuotaSpec {
    #[serde(default)]
    daily_token_limit: Option<i64>,
    #[serde(default)]
    daily_cost_limit: Option<f64>,
}

#[derive(Debug, Deserialize)]
struct RoleSpec {
    email: String,
    role: String,
}

#[derive(Debug, Deserialize)]
struct UsageSpec {
    #[serde(default)]
    tokens: Option<i64>,
    #[serde(default)]
    cost: Option<f64>,
    #[serde(default)]
    recorded_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Serialize)]
struct CreateTenantResponse {
    tenant: Tenant,
}

#[derive(Debug, Serialize)]
struct CreateProjectResponse {
    project: Project,
}

#[derive(Debug, Serialize)]
struct CreateKeyResponse {
    key: ApiKey,
}

#[derive(Debug, Serialize)]
struct ErrorResponse {
    error: String,
}

#[derive(Debug, Serialize)]
struct QuotaResponse {
    quota: TenantQuota,
}

#[derive(Debug, Serialize)]
struct RolesResponse {
    roles: Vec<RoleBinding>,
}

#[derive(Debug, Serialize, Clone)]
struct RoleBinding {
    email: String,
    role: String,
}

#[derive(Debug, Serialize)]
struct UsageResponse {
    records: Vec<BillingRecord>,
}

pub async fn serve(config: ControlPlaneConfig) -> Result<()> {
    let listener = TcpListener::bind(config.addr)
        .await
        .with_context(|| format!("failed to bind control plane on {}", config.addr))?;
    info!(addr = %config.addr, "starting managed control plane");

    let state = Arc::new(ControlPlaneState::default());
    let app = Router::new()
        .route("/healthz", get(healthz))
        .route("/v1/tenants", post(create_tenant).get(list_tenants))
        .route("/v1/tenants/:tenant_id/projects", post(create_project))
        .route(
            "/v1/tenants/:tenant_id/quotas",
            post(upsert_quota).get(get_quota),
        )
        .route(
            "/v1/tenants/:tenant_id/roles",
            post(assign_role).get(list_roles),
        )
        .route(
            "/v1/tenants/:tenant_id/usage",
            post(record_usage).get(list_usage),
        )
        .route(
            "/v1/projects/:project_id/keys",
            post(issue_key).get(list_keys),
        )
        .with_state(state)
        .layer(TraceLayer::new_for_http());

    axum::serve(listener, app)
        .await
        .context("control plane server exited unexpectedly")
}

async fn healthz() -> &'static str {
    "ok"
}

#[instrument(skip(state, payload))]
async fn create_tenant(
    State(state): State<Arc<ControlPlaneState>>,
    Json(payload): Json<TenantSpec>,
) -> impl IntoResponse {
    if payload.name.trim().is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "name is required".into(),
            }),
        );
    }

    let tenant = Tenant {
        id: Uuid::new_v4(),
        name: payload.name.clone(),
        contact: payload.contact.clone(),
        status: "provisioning",
    };

    state.tenants.write().insert(tenant.id, tenant.clone());
    info!(tenant_id = %tenant.id, "provisioned tenant");
    (StatusCode::CREATED, Json(CreateTenantResponse { tenant }))
}

async fn list_tenants(State(state): State<Arc<ControlPlaneState>>) -> Json<Vec<Tenant>> {
    let tenants = state.tenants.read();
    Json(tenants.values().cloned().collect())
}

#[instrument(skip(state, payload))]
async fn create_project(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
    Json(payload): Json<ProjectSpec>,
) -> impl IntoResponse {
    if payload.name.trim().is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "name is required".into(),
            }),
        );
    }

    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let project = Project {
        id: Uuid::new_v4(),
        tenant_id,
        name: payload.name.clone(),
    };
    state.projects.write().insert(project.id, project.clone());
    info!(tenant_id = %tenant_id, project_id = %project.id, "created project");
    (StatusCode::CREATED, Json(CreateProjectResponse { project }))
}

#[instrument(skip(state))]
async fn issue_key(
    State(state): State<Arc<ControlPlaneState>>,
    Path(project_id): Path<Uuid>,
) -> impl IntoResponse {
    let projects = state.projects.read();
    if !projects.contains_key(&project_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "project not found".into(),
            }),
        );
    }
    drop(projects);

    let api_key = ApiKey {
        id: Uuid::new_v4(),
        project_id,
        key: format!("ffk_{}", Uuid::new_v4().without_hyphens()),
    };
    state.keys.write().insert(api_key.id, api_key.clone());
    info!(project_id = %project_id, key_id = %api_key.id, "issued project key");
    (
        StatusCode::CREATED,
        Json(CreateKeyResponse { key: api_key }),
    )
}

async fn list_keys(
    State(state): State<Arc<ControlPlaneState>>,
    Path(project_id): Path<Uuid>,
) -> impl IntoResponse {
    let projects = state.projects.read();
    if !projects.contains_key(&project_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "project not found".into(),
            }),
        );
    }
    drop(projects);

    let keys = state
        .keys
        .read()
        .values()
        .filter(|key| key.project_id == project_id)
        .cloned()
        .collect::<Vec<_>>();

    Json(keys)
}

#[instrument(skip(state, payload))]
async fn assign_role(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
    Json(payload): Json<RoleSpec>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let role = payload.role.to_ascii_lowercase();
    if !VALID_ROLES.contains(&role.as_str()) {
        return (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: format!("unknown role '{}'", payload.role),
            }),
        );
    }
    if payload.email.trim().is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "email is required".into(),
            }),
        );
    }

    let mut roles = state.tenant_roles.write();
    roles
        .entry(tenant_id)
        .or_default()
        .insert(payload.email.clone(), role.clone());
    info!(tenant_id = %tenant_id, email = %payload.email, role = %role, "assigned tenant role");
    (
        StatusCode::OK,
        Json(RolesResponse {
            roles: roles[&tenant_id]
                .iter()
                .map(|(email, role)| RoleBinding {
                    email: email.clone(),
                    role: role.clone(),
                })
                .collect(),
        }),
    )
}

async fn list_roles(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let roles = state.tenant_roles.read();
    let bindings = roles
        .get(&tenant_id)
        .map(|entries| {
            entries
                .iter()
                .map(|(email, role)| RoleBinding {
                    email: email.clone(),
                    role: role.clone(),
                })
                .collect::<Vec<_>>()
        })
        .unwrap_or_default();
    Json(RolesResponse { roles: bindings })
}

#[instrument(skip(state, payload))]
async fn record_usage(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
    Json(payload): Json<UsageSpec>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let tokens = payload.tokens.unwrap_or_default().max(0);
    let cost = payload.cost.unwrap_or_default().max(0.0);
    let record = BillingRecord {
        tenant_id,
        recorded_at: payload.recorded_at.unwrap_or_else(Utc::now),
        tokens,
        cost,
    };
    state.billing.write().push(record.clone());
    info!(tenant_id = %tenant_id, tokens = tokens, cost = cost, "recorded usage sample");
    (StatusCode::CREATED, Json(record))
}

async fn list_usage(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let records = state
        .billing
        .read()
        .iter()
        .filter(|record| record.tenant_id == tenant_id)
        .cloned()
        .collect::<Vec<_>>();
    Json(UsageResponse { records })
}

#[instrument(skip(state, payload))]
async fn upsert_quota(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
    Json(payload): Json<QuotaSpec>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    if payload.daily_token_limit.is_none() && payload.daily_cost_limit.is_none() {
        state.quotas.write().remove(&tenant_id);
        info!(tenant_id = %tenant_id, "cleared tenant quota limits");
        return StatusCode::NO_CONTENT;
    }

    let quota = TenantQuota {
        tenant_id,
        daily_token_limit: payload.daily_token_limit,
        daily_cost_limit: payload.daily_cost_limit,
    };
    state.quotas.write().insert(tenant_id, quota.clone());
    info!(tenant_id = %tenant_id, "updated tenant quota limits");
    (StatusCode::OK, Json(QuotaResponse { quota }))
}

async fn get_quota(
    State(state): State<Arc<ControlPlaneState>>,
    Path(tenant_id): Path<Uuid>,
) -> impl IntoResponse {
    let tenants = state.tenants.read();
    if !tenants.contains_key(&tenant_id) {
        return (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "tenant not found".into(),
            }),
        );
    }
    drop(tenants);

    let quotas = state.quotas.read();
    match quotas.get(&tenant_id) {
        Some(quota) => (
            StatusCode::OK,
            Json(QuotaResponse {
                quota: quota.clone(),
            }),
        ),
        None => (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "quota not configured".into(),
            }),
        ),
    }
}

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

    fn seeded_state() -> (Arc<ControlPlaneState>, Uuid) {
        let state = ControlPlaneState::default();
        let tenant = Tenant {
            id: Uuid::new_v4(),
            name: "Acme".into(),
            contact: Some("ops@acme.test".into()),
            status: "active",
        };
        let id = tenant.id;
        state.tenants.write().insert(id, tenant);
        (Arc::new(state), id)
    }

    #[tokio::test]
    async fn assign_role_and_list_roles() {
        let (state, tenant_id) = seeded_state();

        let response = assign_role(
            State(state.clone()),
            Path(tenant_id),
            Json(RoleSpec {
                email: "owner@acme.test".into(),
                role: "admin".into(),
            }),
        )
        .await;

        assert_eq!(response.0, StatusCode::OK);

        let Json(listed) = list_roles(State(state.clone()), Path(tenant_id)).await;
        assert_eq!(listed.roles.len(), 1);
        assert_eq!(listed.roles[0].role, "admin");
    }

    #[tokio::test]
    async fn assign_role_rejects_unknown_role() {
        let (state, tenant_id) = seeded_state();

        let response = assign_role(
            State(state),
            Path(tenant_id),
            Json(RoleSpec {
                email: "owner@acme.test".into(),
                role: "unknown".into(),
            }),
        )
        .await;

        assert_eq!(response.0, StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn record_and_list_usage() {
        let (state, tenant_id) = seeded_state();

        let response = record_usage(
            State(state.clone()),
            Path(tenant_id),
            Json(UsageSpec {
                tokens: Some(120),
                cost: Some(0.45),
                recorded_at: None,
            }),
        )
        .await;

        assert_eq!(response.0, StatusCode::CREATED);

        let Json(listed) = list_usage(State(state), Path(tenant_id)).await;
        assert_eq!(listed.records.len(), 1);
        assert_eq!(listed.records[0].tokens, 120);
    }
}