routing.rs

//! Shared routing pipeline for agent tunnel.
//!
//! Used by both connection forwarding (`fwd.rs`) and KDC proxy (`kdc_proxy.rs`)
//! to ensure consistent routing behavior and error messages.

use std::net::IpAddr;
use std::sync::Arc;

use anyhow::{Result, anyhow};
use uuid::Uuid;

use super::listener::AgentTunnelHandle;
use super::registry::{AgentPeer, AgentRegistry};
use super::stream::TunnelStream;

/// Result of the routing pipeline.
///
/// Each variant carries enough context for the caller to produce an actionable error message.
#[derive(Debug)]
pub enum RoutingDecision {
    /// Route through these agent candidates (try in order, first success wins).
    ViaAgent(Vec<Arc<AgentPeer>>),
    /// Explicit agent_id was specified but not found in registry.
    ExplicitAgentNotFound(Uuid),
    /// No agent matched — caller should attempt direct connection.
    Direct,
}

/// Determines how to route a connection to the given target.
///
/// Pipeline (in order of priority):
/// 1. Explicit agent_id (from JWT) → route to that agent
/// 2. IP target → subnet match against agent advertisements
/// 3. Hostname target → domain suffix match (longest wins)
/// 4. No match → direct connection
pub fn resolve_route(registry: &AgentRegistry, explicit_agent_id: Option<Uuid>, target_host: &str) -> RoutingDecision {
    // Step 1: Explicit agent ID (from JWT)
    if let Some(agent_id) = explicit_agent_id {
        if let Some(agent) = registry.get(&agent_id) {
            return RoutingDecision::ViaAgent(vec![agent]);
        }
        return RoutingDecision::ExplicitAgentNotFound(agent_id);
    }

    // Step 2: Target is an IP address → subnet match
    if let Ok(ip) = target_host.parse::<IpAddr>() {
        let agents = registry.find_agents_for_target(ip);
        if !agents.is_empty() {
            return RoutingDecision::ViaAgent(agents);
        }
        return RoutingDecision::Direct;
    }

    // Step 3: Target is a hostname → domain suffix match (longest wins)
    let agents = registry.select_agents_for_domain(target_host);
    if !agents.is_empty() {
        return RoutingDecision::ViaAgent(agents);
    }

    // Step 4: No match → direct connect
    RoutingDecision::Direct
}

/// Try connecting to target through agent candidates (try-fail-retry).
///
/// Returns the connected `TunnelStream` and the agent that succeeded.
///
/// Callers must handle `RoutingDecision::ExplicitAgentNotFound` and
/// `RoutingDecision::Direct` before calling this function.
pub async fn route_and_connect(
    handle: &AgentTunnelHandle,
    candidates: &[Arc<AgentPeer>],
    session_id: Uuid,
    target: &str,
) -> Result<(TunnelStream, Arc<AgentPeer>)> {
    assert!(!candidates.is_empty(), "route_and_connect called with empty candidates");

    let mut last_error = None;

    for agent in candidates {
        info!(
            agent_id = %agent.agent_id,
            agent_name = %agent.name,
            %target,
            "Routing via agent tunnel"
        );

        match handle.connect_via_agent(agent.agent_id, session_id, target).await {
            Ok(stream) => {
                info!(
                    agent_id = %agent.agent_id,
                    agent_name = %agent.name,
                    %target,
                    "Agent tunnel connection established"
                );
                return Ok((stream, Arc::clone(agent)));
            }
            Err(error) => {
                warn!(
                    agent_id = %agent.agent_id,
                    agent_name = %agent.name,
                    %target,
                    error = format!("{error:#}"),
                    "Agent tunnel connection failed, trying next candidate"
                );
                last_error = Some(error);
            }
        }
    }

    let agent_names: Vec<&str> = candidates.iter().map(|a| a.name.as_str()).collect();
    let last_err_msg = last_error.as_ref().map(|e| format!("{e:#}")).unwrap_or_default();

    error!(
        agent_count = candidates.len(),
        %target,
        agents = ?agent_names,
        last_error = %last_err_msg,
        "All agent tunnel candidates failed"
    );

    Err(last_error.unwrap_or_else(|| {
        anyhow!(
            "All {} agents matching target '{}' failed to connect. Agents tried: [{}]",
            candidates.len(),
            target,
            agent_names.join(", "),
        )
    }))
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::Ordering;

    use agent_tunnel_proto::DomainAdvertisement;

    use super::*;
    use crate::agent_tunnel::registry::AgentPeer;

    fn make_peer(name: &str) -> Arc<AgentPeer> {
        Arc::new(AgentPeer::new(
            Uuid::new_v4(),
            name.to_owned(),
            "sha256:test".to_owned(),
        ))
    }

    fn domain(name: &str) -> DomainAdvertisement {
        DomainAdvertisement {
            domain: name.to_owned(),
            auto_detected: false,
        }
    }

    #[test]
    fn route_explicit_agent_id() {
        let registry = AgentRegistry::new();
        let peer = make_peer("agent-a");
        let agent_id = peer.agent_id;
        registry.register(Arc::clone(&peer));

        match resolve_route(&registry, Some(agent_id), "anything") {
            RoutingDecision::ViaAgent(agents) => {
                assert_eq!(agents.len(), 1);
                assert_eq!(agents[0].agent_id, agent_id);
            }
            other => panic!("expected ViaAgent, got {other:?}"),
        }
    }

    #[test]
    fn route_explicit_agent_id_not_found() {
        let registry = AgentRegistry::new();
        let bogus_id = Uuid::new_v4();

        match resolve_route(&registry, Some(bogus_id), "anything") {
            RoutingDecision::ExplicitAgentNotFound(id) => {
                assert_eq!(id, bogus_id);
            }
            other => panic!("expected ExplicitAgentNotFound, got {other:?}"),
        }
    }

    #[test]
    fn route_ip_target_via_subnet() {
        let registry = AgentRegistry::new();
        let peer = make_peer("agent-a");
        let agent_id = peer.agent_id;
        let subnet: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer.update_routes(1, vec![subnet], vec![]);
        registry.register(peer);

        match resolve_route(&registry, None, "10.1.5.50") {
            RoutingDecision::ViaAgent(agents) => {
                assert_eq!(agents[0].agent_id, agent_id);
            }
            other => panic!("expected ViaAgent, got {other:?}"),
        }
    }

    #[test]
    fn route_hostname_via_domain() {
        let registry = AgentRegistry::new();
        let peer = make_peer("agent-a");
        let agent_id = peer.agent_id;
        let subnet: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer.update_routes(1, vec![subnet], vec![domain("contoso.local")]);
        registry.register(peer);

        match resolve_route(&registry, None, "dc01.contoso.local") {
            RoutingDecision::ViaAgent(agents) => {
                assert_eq!(agents[0].agent_id, agent_id);
            }
            other => panic!("expected ViaAgent, got {other:?}"),
        }
    }

    #[test]
    fn route_no_match_returns_direct() {
        let registry = AgentRegistry::new();
        let peer = make_peer("agent-a");
        let subnet: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer.update_routes(1, vec![subnet], vec![domain("contoso.local")]);
        registry.register(peer);

        assert!(matches!(
            resolve_route(&registry, None, "external.example.com"),
            RoutingDecision::Direct
        ));
    }

    #[test]
    fn route_ip_no_match_returns_direct() {
        let registry = AgentRegistry::new();
        let peer = make_peer("agent-a");
        let subnet: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer.update_routes(1, vec![subnet], vec![]);
        registry.register(peer);

        assert!(matches!(
            resolve_route(&registry, None, "192.168.1.1"),
            RoutingDecision::Direct
        ));
    }

    #[test]
    fn route_skips_offline_agents() {
        let registry = AgentRegistry::new();
        let peer = make_peer("offline-agent");
        let subnet: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer.update_routes(1, vec![subnet], vec![domain("contoso.local")]);
        peer.last_seen.store(0, Ordering::Release);
        registry.register(peer);

        assert!(matches!(
            resolve_route(&registry, None, "dc01.contoso.local"),
            RoutingDecision::Direct
        ));
    }

    #[test]
    fn route_domain_match_returns_multiple_agents_ordered() {
        let registry = AgentRegistry::new();

        let peer_a = make_peer("agent-a");
        let subnet_a: ipnetwork::Ipv4Network = "10.1.0.0/16".parse().expect("valid test subnet");
        peer_a.update_routes(1, vec![subnet_a], vec![domain("contoso.local")]);
        registry.register(Arc::clone(&peer_a));

        std::thread::sleep(std::time::Duration::from_millis(10));

        let peer_b = make_peer("agent-b");
        let id_b = peer_b.agent_id;
        let subnet_b: ipnetwork::Ipv4Network = "10.2.0.0/16".parse().expect("valid test subnet");
        peer_b.update_routes(1, vec![subnet_b], vec![domain("contoso.local")]);
        registry.register(Arc::clone(&peer_b));

        match resolve_route(&registry, None, "dc01.contoso.local") {
            RoutingDecision::ViaAgent(agents) => {
                assert_eq!(agents.len(), 2);
                assert_eq!(agents[0].agent_id, id_b, "most recent first");
            }
            other => panic!("expected ViaAgent, got {other:?}"),
        }
    }
}