gitdataai/apps/gingress/src/controller/ingress_watcher.rs

//! Watches Kubernetes Ingress resources and converts them to routing rules.

use futures::pin_mut;
use futures::StreamExt;
use gingress_proxy::config::{
    ConfigStore, HeaderOp, PathType, RateLimitPolicy, RouteRule, SessionAffinityConfig,
};
use k8s_openapi::api::networking::v1::{HTTPIngressPath, Ingress};
use kube::ResourceExt;
use kube::runtime::watcher::{self, Event};
use std::collections::BTreeMap;
use std::sync::Arc;

/// Watch Ingress resources and update the ConfigStore.
///
/// After each event, the `on_change` callback is invoked so the reconciler
/// can cross-reference all fragments into a complete ProxyConfig.
pub async fn watch_ingresses(
    client: Arc<kube::Client>,
    store: Arc<ConfigStore>,
    ingress_class: String,
    namespace: Option<String>,
    on_change: Arc<dyn Fn() + Send + Sync>,
) {
    let api = kube::Api::<Ingress>::all(client.as_ref().clone());

    let config = watcher::Config {
        field_selector: namespace.as_ref().map(|ns| format!("metadata.namespace={}", ns)),
        ..Default::default()
    };

    let ingress_watcher = watcher::watcher(api, config);
    pin_mut!(ingress_watcher);

    while let Some(event) = ingress_watcher.next().await {
        match event {
            Ok(Event::Apply(ingress)) => {
                let name = ingress.name_any();
                let ns = ingress.namespace().unwrap_or_default();
                if is_gingress_class(&ingress, &ingress_class) {
                    process_ingress(&ingress, &store, &ingress_class);
                    on_change();
                    tracing::info!(namespace = %ns, name = %name, "Ingress applied");
                }
            }
            Ok(Event::Init) => {
                store.remove_prefix("ingress:");
                store.remove_prefix("tls-host:");
                tracing::info!("Ingress watcher re-initializing");
            }
            Ok(Event::InitApply(ingress)) => {
                if is_gingress_class(&ingress, &ingress_class) {
                    process_ingress(&ingress, &store, &ingress_class);
                }
            }
            Ok(Event::InitDone) => {
                store.signal_reload();
                on_change();
                tracing::info!("Ingress watcher init complete");
            }
            Ok(Event::Delete(ingress)) => {
                if is_gingress_class(&ingress, &ingress_class) {
                    remove_ingress_routes(&ingress, &store);
                    on_change();
                    tracing::info!(
                        name = %ingress.name_any(),
                        namespace = %ingress.namespace().unwrap_or_default(),
                        "Ingress deleted, routes removed"
                    );
                }
            }
            Err(e) => {
                tracing::error!("Ingress watcher error: {}", e);
            }
        }
    }
}

/// Check if an Ingress specifies the gingress class.
fn is_gingress_class(ingress: &Ingress, class_name: &str) -> bool {
    ingress
        .spec
        .as_ref()
        .and_then(|s| s.ingress_class_name.as_deref())
        == Some(class_name)
}

/// Process an Ingress resource: extract routes and update the store.
fn process_ingress(ingress: &Ingress, store: &ConfigStore, _ingress_class: &str) {
    let namespace = ingress.namespace().unwrap_or_default();
    let name = ingress.name_any();
    let spec = match ingress.spec.as_ref() {
        Some(s) => s,
        None => return,
    };

    // Build an ingress-scoped prefix so we can clean up old routes for this Ingress
    let ingress_prefix = format!("ingress:{}/{}:", namespace, name);

    // Remove old route entries scoped to this Ingress
    let old_route_keys = store.keys_with_prefix(&format!("{}route:", ingress_prefix));
    for key in &old_route_keys {
        store.remove(key);
    }

    // Process routing rules
    if let Some(rules) = &spec.rules {
        for rule in rules {
            let host = rule.host.as_deref().unwrap_or("*");
            if let Some(http) = &rule.http {
                let mut routes: Vec<RouteRule> = Vec::new();
                for path_item in &http.paths {
                    routes.push(ingress_path_to_route(host, path_item, &namespace));
                }
                // Store per-ingress routes so we can clean up on delete
                let route_key = format!("{}route:{}", ingress_prefix, host);
                store.set(&route_key, &routes);
            }
        }
    }

    // Process TLS: map secretName -> hosts so the reconciler can cross-reference
    if let Some(tls_entries) = &spec.tls {
        for tls in tls_entries {
            let secret_name = tls.secret_name.as_deref().unwrap_or_default();
            let hosts: Vec<String> = tls.hosts.clone().unwrap_or_default();
            let tls_host_key = format!("tls-host:{}", secret_name);
            store.set(&tls_host_key, &hosts);
        }
    }

    // Process annotations for advanced features
    let annotations = ingress.annotations();
    process_annotations(&annotations, &ingress_prefix, store);

    store.signal_reload();
}

/// Convert a Kubernetes Ingress path to an internal RouteRule.
fn ingress_path_to_route(host: &str, path: &HTTPIngressPath, namespace: &str) -> RouteRule {
    let service = path.backend.service.as_ref()
        .expect("Ingress backend must reference a service");

    RouteRule {
        host: host.to_string(),
        path: path.path.clone().unwrap_or_else(|| "/".to_string()),
        path_type: match path.path_type.as_str() {
            "Prefix" => PathType::Prefix,
            "Exact" => PathType::Exact,
            _ => PathType::ImplementationSpecific,
        },
        backend: gingress_proxy::config::Backend {
            namespace: namespace.to_string(),
            name: service.name.clone(),
            port: service.port.as_ref().and_then(|p| p.number).unwrap_or(80) as u16,
        },
    }
}

/// Annotation keys for GIngress features.
const ANN_RATE_LIMIT: &str = "gingress.io/rate-limit";
const ANN_RATE_LIMIT_BURST: &str = "gingress.io/rate-limit-burst";
const ANN_REQUEST_HEADERS: &str = "gingress.io/request-headers";
const ANN_WEBSOCKET: &str = "gingress.io/websocket";
const ANN_SESSION_AFFINITY: &str = "gingress.io/session-affinity";

/// Parse Ingress annotations and write corresponding ConfigStore entries.
///
/// Supported annotations:
/// - `gingress.io/rate-limit` — "RPS" or "RPS/BURST" (e.g., "100" or "100/200")
/// - `gingress.io/rate-limit-burst` — Override burst size
/// - `gingress.io/request-headers` — JSON array of header operations
/// - `gingress.io/websocket` — "true" to enable WebSocket upgrade for this host
/// - `gingress.io/session-affinity` — "cookie" or "cookie:NAME:TTL_SECONDS"
fn process_annotations(
    annotations: &BTreeMap<String, String>,
    ingress_prefix: &str,
    store: &ConfigStore,
) {
    // Collect hosts from the ingress routes that were just stored
    let route_keys = store.keys_with_prefix(&format!("{}route:", ingress_prefix));
    let hosts: Vec<String> = route_keys
        .iter()
        .filter_map(|k| k.split(":route:").nth(1).map(String::from))
        .collect();

    if hosts.is_empty() {
        return;
    }

    // Remove old per-host annotation keys (handles annotation removal/update)
    for host in &hosts {
        store.remove(&format!("rate_limit:{}", host));
        store.remove(&format!("headers:{}", host));
        store.remove(&format!("session_affinity:{}", host));
    }

    // Remove this ingress's hosts from the global websocket list
    prune_websocket_hosts(store, &hosts);

    // ── Rate limiting ──
    if let Some(val) = annotations.get(ANN_RATE_LIMIT) {
        let (rps, burst) = parse_rate_limit(val, annotations.get(ANN_RATE_LIMIT_BURST));
        for host in &hosts {
            store.set(
                &format!("rate_limit:{}", host),
                &RateLimitPolicy {
                    host: host.clone(),
                    requests_per_second: rps,
                    burst_size: burst,
                },
            );
        }
    }

    // ── Header operations (request) ──
    if let Some(val) = annotations.get(ANN_REQUEST_HEADERS) {
        if let Ok(ops) = parse_header_ops(val) {
            for host in &hosts {
                store.set(&format!("headers:{}", host), &ops);
            }
        } else {
            tracing::warn!(annotation = %ANN_REQUEST_HEADERS, value = %val, "Invalid header ops JSON");
        }
    }

    // ── WebSocket ──
    if let Some(val) = annotations.get(ANN_WEBSOCKET) {
        if val.trim().to_lowercase() == "true" {
            let mut ws_hosts: Vec<String> = hosts.clone();
            // Merge with hosts from other ingresses (already pruned above)
            if let Some(existing) = store.get::<Vec<String>>("websocket:hosts") {
                for h in existing {
                    if !ws_hosts.contains(&h) {
                        ws_hosts.push(h);
                    }
                }
            }
            store.set("websocket:hosts", &ws_hosts);
        }
    }

    // ── Session affinity ──
    if let Some(val) = annotations.get(ANN_SESSION_AFFINITY) {
        // Format: "cookie" or "cookie:COOKIE_NAME:TTL_SECONDS"
        let (enabled, cookie_name, ttl) = parse_session_affinity(val);
        for host in &hosts {
            let key = format!("session_affinity:{}", host);
            store.set(
                &key,
                &SessionAffinityConfig {
                    enabled,
                    cookie_name: cookie_name.clone(),
                    cookie_ttl_seconds: ttl,
                },
            );
        }
    }
}

fn parse_rate_limit(val: &str, burst_override: Option<&String>) -> (u32, u32) {
    let val = val.trim();
    if let Some((rps_str, burst_str)) = val.split_once('/') {
        let rps = rps_str.parse().unwrap_or(0);
        let burst = burst_str.parse().unwrap_or(rps);
        (rps, burst)
    } else {
        let rps = val.parse().unwrap_or(0);
        let burst = burst_override
            .and_then(|b| b.parse().ok())
            .unwrap_or(rps * 2);
        (rps, burst)
    }
}

#[derive(serde::Deserialize)]
struct HeaderOpAnnotation {
    op: String,
    name: String,
    #[serde(default)]
    value: Option<String>,
}

fn parse_header_ops(val: &str) -> anyhow::Result<Vec<HeaderOp>> {
    let items: Vec<HeaderOpAnnotation> = serde_json::from_str(val)?;
    items
        .into_iter()
        .map(|item| {
            Ok(match item.op.as_str() {
                "set" => HeaderOp::Set {
                    name: item.name,
                    value: item.value.unwrap_or_default(),
                },
                "add" => HeaderOp::Add {
                    name: item.name,
                    value: item.value.unwrap_or_default(),
                },
                "remove" => HeaderOp::Remove { name: item.name },
                _ => anyhow::bail!("Unknown header op: {}", item.op),
            })
        })
        .collect()
}

fn parse_session_affinity(val: &str) -> (bool, String, u64) {
    let val = val.trim();
    if val.eq_ignore_ascii_case("cookie") || val.eq_ignore_ascii_case("true") {
        return (true, "GINGRESS_AFFINITY".into(), 3600);
    }
    // Format: "cookie:COOKIE_NAME:TTL"
    let parts: Vec<&str> = val.split(':').collect();
    if parts.len() >= 3 {
        let name = parts[1].to_string();
        let ttl = parts[2].parse().unwrap_or(3600);
        (true, name, ttl)
    } else if parts.len() == 2 {
        let name = parts[1].to_string();
        (true, name, 3600)
    } else {
        (false, String::new(), 0)
    }
}

/// Remove a set of hosts from the global websocket host list (scoped cleanup).
fn prune_websocket_hosts(store: &ConfigStore, hosts_to_remove: &[String]) {
    if let Some(mut existing) = store.get::<Vec<String>>("websocket:hosts") {
        existing.retain(|h| !hosts_to_remove.contains(h));
        if existing.is_empty() {
            store.remove("websocket:hosts");
        } else {
            store.set("websocket:hosts", &existing);
        }
    }
}

/// Remove all routes associated with a deleted Ingress.
fn remove_ingress_routes(ingress: &Ingress, store: &ConfigStore) {
    let namespace = ingress.namespace().unwrap_or_default();
    let name = ingress.name_any();
    let ingress_prefix = format!("ingress:{}/{}:", namespace, name);

    // Collect hosts before deleting routes so we can clean up per-host annotation keys
    let host_keys: Vec<String> = store
        .keys_with_prefix(&format!("{}route:", ingress_prefix))
        .iter()
        .filter_map(|k| k.split(":route:").nth(1).map(String::from))
        .collect();

    // Remove all route entries for this Ingress
    store.remove_prefix(&ingress_prefix);

    // Remove per-host annotation-derived keys
    for host in &host_keys {
        store.remove(&format!("rate_limit:{}", host));
        store.remove(&format!("headers:{}", host));
        store.remove(&format!("session_affinity:{}", host));
    }
    // Scoped: only remove this ingress's hosts from the global websocket list
    prune_websocket_hosts(store, &host_keys);

    // Remove TLS host mappings
    if let Some(spec) = ingress.spec.as_ref() {
        if let Some(tls_entries) = &spec.tls {
            for tls in tls_entries {
                let sn = tls.secret_name.as_deref().unwrap_or_default();
                store.remove(&format!("tls-host:{}", sn));
            }
        }
    }

    store.signal_reload();
}