featherChat/warzone/crates/warzone-server/src/routes/keys.rs

use axum::{
    extract::{Path, State},
    routing::{get, post},
    Json, Router,
};
use serde::{Deserialize, Serialize};

use crate::state::AppState;

pub fn routes() -> Router<AppState> {
    Router::new()
        .route("/keys/register", post(register_keys))
        .route("/keys/replenish", post(replenish_otpks))
        .route("/keys/list", get(list_keys))
        .route("/keys/:fingerprint", get(get_bundle))
        .route("/keys/:fingerprint/otpk-count", get(otpk_count))
        .route("/keys/:fingerprint/devices", get(list_devices))
}

/// Debug endpoint: list all registered fingerprints.
async fn list_keys(State(state): State<AppState>) -> Json<serde_json::Value> {
    let keys: Vec<String> = state
        .db
        .keys
        .iter()
        .filter_map(|item| {
            item.ok()
                .and_then(|(k, _)| String::from_utf8(k.to_vec()).ok())
        })
        .collect();
    tracing::info!("Listed {} registered keys", keys.len());
    Json(serde_json::json!({ "keys": keys, "count": keys.len() }))
}

/// Normalize fingerprint: strip colons, lowercase.
fn normalize_fp(fp: &str) -> String {
    fp.chars()
        .filter(|c| c.is_ascii_hexdigit())
        .collect::<String>()
        .to_lowercase()
}

#[derive(Deserialize)]
struct RegisterRequest {
    fingerprint: String,
    #[serde(default)]
    device_id: Option<String>,
    bundle: Vec<u8>,
    #[serde(default)]
    eth_address: Option<String>,
}

#[derive(Serialize)]
struct RegisterResponse {
    ok: bool,
}

async fn register_keys(

    State(state): State<AppState>,
    Json(req): Json<RegisterRequest>,
) -> Json<RegisterResponse> {
    let fp = normalize_fp(&req.fingerprint);
    let device_id = req.device_id.unwrap_or_else(|| "default".to_string());

    // Store bundle keyed by fingerprint (primary, used for lookup)
    let _ = state.db.keys.insert(fp.as_bytes(), req.bundle.clone());

    // Also store per-device: device:<fp>:<device_id> → bundle
    let device_key = format!("device:{}:{}", fp, device_id);
    let _ = state.db.keys.insert(device_key.as_bytes(), req.bundle);

    // Store ETH address mapping if provided
    if let Some(ref eth) = req.eth_address {
        let eth_lower = eth.to_lowercase();
        // eth -> fp
        let _ = state.db.eth_addresses.insert(eth_lower.as_bytes(), fp.as_bytes());
        // fp -> eth (reverse lookup)
        let _ = state.db.eth_addresses.insert(format!("rev:{}", fp).as_bytes(), eth_lower.as_bytes());
        tracing::info!("ETH address mapped: {} -> {}", eth_lower, fp);
    }

    tracing::info!("Registered bundle for {} (device: {})", fp, device_id);
    Json(RegisterResponse { ok: true })
}

async fn get_bundle(
    State(state): State<AppState>,
    Path(fingerprint): Path<String>,
) -> Result<Json<serde_json::Value>, axum::http::StatusCode> {
    let key = normalize_fp(&fingerprint);
    tracing::info!("get_bundle: raw path='{}', normalized='{}'", fingerprint, key);

    // Debug: list what's in the DB
    let all_keys: Vec<String> = state.db.keys.iter()
        .filter_map(|r| r.ok().and_then(|(k, _)| String::from_utf8(k.to_vec()).ok()))
        .collect();
    tracing::info!("get_bundle: DB contains {} keys: {:?}", all_keys.len(), all_keys);

    // Check if this fingerprint registered locally (has a device: entry)
    let device_prefix = format!("device:{}:", key);
    let is_local = state.db.keys.scan_prefix(device_prefix.as_bytes()).next().is_some();

    // For remote clients, always proxy from the federation peer (bundles may change)
    if !is_local {
        if let Some(ref federation) = state.federation {
            if let Some(bundle_bytes) = federation.fetch_remote_bundle(&key).await {
                tracing::info!("get_bundle: PROXIED from federation peer for {}", key);
                return Ok(Json(serde_json::json!({
                    "fingerprint": fingerprint,
                    "bundle": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bundle_bytes),
                })));
            }
        }
    }

    match state.db.keys.get(key.as_bytes()) {
        Ok(Some(data)) => {
            tracing::info!("get_bundle: FOUND {} bytes for {} (local={})", data.len(), key, is_local);
            Ok(Json(serde_json::json!({
                "fingerprint": fingerprint,
                "bundle": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &data),
            })))
        }
        Ok(None) => {
            tracing::warn!("get_bundle: NOT FOUND for key '{}'", key);
            Err(axum::http::StatusCode::NOT_FOUND)
        }
        Err(e) => {
            tracing::error!("get_bundle: DB error: {}", e);
            Err(axum::http::StatusCode::INTERNAL_SERVER_ERROR)
        }
    }
}

/// Check how many one-time pre-keys remain for a fingerprint.
async fn otpk_count(
    State(state): State<AppState>,
    Path(fingerprint): Path<String>,
) -> Json<serde_json::Value> {
    let fp = normalize_fp(&fingerprint);
    let prefix = format!("otpk:{}:", fp);
    let count = state.db.keys.scan_prefix(prefix.as_bytes()).count();
    Json(serde_json::json!({ "fingerprint": fp, "otpk_count": count }))
}

#[derive(Deserialize)]
struct ReplenishRequest {
    fingerprint: String,
    /// One-time pre-keys: list of {id, public_key_hex}
    otpks: Vec<OtpkEntry>,
}

#[derive(Deserialize)]
struct OtpkEntry {
    id: u32,
    public_key: String, // hex-encoded 32-byte X25519 public key
}

/// Upload additional one-time pre-keys.
async fn replenish_otpks(

    State(state): State<AppState>,
    Json(req): Json<ReplenishRequest>,
) -> Json<serde_json::Value> {
    let fp = normalize_fp(&req.fingerprint);
    let mut stored = 0;

    for otpk in &req.otpks {
        let key = format!("otpk:{}:{}", fp, otpk.id);
        let _ = state.db.keys.insert(key.as_bytes(), otpk.public_key.as_bytes());
        stored += 1;
    }

    let prefix = format!("otpk:{}:", fp);
    let total = state.db.keys.scan_prefix(prefix.as_bytes()).count();

    tracing::info!("Replenished {} OTPKs for {} (total: {})", stored, fp, total);
    Json(serde_json::json!({ "ok": true, "stored": stored, "total": total }))
}

/// List all registered devices for a fingerprint.
async fn list_devices(
    State(state): State<AppState>,
    Path(fingerprint): Path<String>,
) -> Json<serde_json::Value> {
    let fp = normalize_fp(&fingerprint);
    let prefix = format!("device:{}:", fp);
    let devices: Vec<String> = state.db.keys.scan_prefix(prefix.as_bytes())
        .filter_map(|item| {
            item.ok().and_then(|(k, _)| {
                let key_str = String::from_utf8_lossy(&k).to_string();
                // key format: device:<fp>:<device_id>
                key_str.rsplit(':').next().map(|s| s.to_string())
            })
        })
        .collect();
    Json(serde_json::json!({ "fingerprint": fp, "devices": devices, "count": devices.len() }))
}