feat: transparent routing through agent tunnel

devolutions-gateway/src/agent_tunnel/mod.rs

@@ -7,8 +7,13 @@ pub mod cert;
 pub mod enrollment_store;
 pub mod listener;
 pub mod registry;
+pub mod routing;
 pub mod stream;
 
+// Integration test needs rewriting for Quinn — kept as local-only file.
+// #[cfg(test)]
+// mod integration_test;
+
 pub use enrollment_store::EnrollmentTokenStore;
 pub use listener::{AgentTunnelHandle, AgentTunnelListener};
 pub use registry::AgentRegistry;

devolutions-gateway/src/agent_tunnel/routing.rs

@@ -0,0 +1,287 @@
+//! 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:?}"),
+        }
+    }
+}