feat(gateway): allow out-of-order allow_access requests (#6403)

Currently, the gateway requires a strict ordering of first receiving a
`request_connection` message, following by multiple `allow_access`
messages. Additionally, access can be granted as part of the initial
`request_connection` message too.

This isn't an ideal design. Setting up a new connection is infallible,
all we need to do is send our ICE credentials back to the client.
However, untangling that will require a bit more effort.

Starting with #6335, following this strict order on the client is a more
difficult. Whilst we can send them in order, it is harder to maintain
those ordering guarantees across all our systems.

To avoid this, we change the gateway to perform an upsert for its local
ACLs for a client. In case that an `allow_access` call would somehow get
to the gateway earlier, we can simply already create the `Peer` and only
set up the actual connection later.

---------

Signed-off-by: Jamil <jamilbk@users.noreply.github.com>
Co-authored-by: Jamil <jamilbk@users.noreply.github.com>

.github/workflows/_integration_tests.yml

@@ -103,6 +103,7 @@ jobs:
           - name: direct-curl-api-restart
           - name: direct-dns-api-down
           - name: direct-dns-relay-down
+          - name: direct-dns-two-resources
           - name: direct-dns
           - name: direct-download-roaming-network
             # Too noisy can cause flaky tests due to the amount of data

elixir/apps/api/lib/api/gateway/channel.ex

@@ -131,6 +131,7 @@ defmodule API.Gateway.Channel do
       } do
       :ok = Flows.subscribe_to_flow_expiration_events(flow_id)
 
+      client = Clients.fetch_client_by_id!(client_id)
       resource = Resources.fetch_resource_by_id!(resource_id)
 
       case API.Client.Channel.map_or_drop_compatible_resource(
@@ -150,7 +151,9 @@ defmodule API.Gateway.Channel do
             flow_id: flow_id,
             resource: Views.Resource.render(resource),
             expires_at: DateTime.to_unix(authorization_expires_at, :second),
-            payload: payload
+            payload: payload,
+            client_ipv4: client.ipv4,
+            client_ipv6: client.ipv6
           })
 
           Logger.debug("Awaiting gateway connection_ready message",

elixir/apps/api/test/api/gateway/channel_test.exs

@@ -122,6 +122,8 @@ defmodule API.Gateway.ChannelTest do
       assert payload.ref
       assert payload.flow_id == flow_id
       assert payload.client_id == client.id
+      assert payload.client_ipv4 == client.ipv4
+      assert payload.client_ipv6 == client.ipv6
       assert DateTime.from_unix!(payload.expires_at) == DateTime.truncate(expires_at, :second)
     end
 

rust/connlib/tunnel/src/gateway.rs

@@ -2,7 +2,7 @@ use crate::peer::ClientOnGateway;
 use crate::peer_store::PeerStore;
 use crate::utils::earliest;
 use crate::{GatewayEvent, GatewayTunnel, BUF_SIZE};
-use anyhow::{bail, Context};
+use anyhow::bail;
 use boringtun::x25519::PublicKey;
 use chrono::{DateTime, Utc};
 use connlib_shared::messages::{
@@ -37,19 +37,13 @@ impl GatewayTunnel {
     }
 
     /// Accept a connection request from a client.
-    #[allow(clippy::too_many_arguments)]
     pub fn accept(
         &mut self,
         client_id: ClientId,
         key: Secret<Key>,
         offer: Offer,
         client: PublicKey,
-        ipv4: Ipv4Addr,
-        ipv6: Ipv6Addr,
-        domain: Option<(DomainName, Vec<IpAddr>)>,
-        expires_at: Option<DateTime<Utc>>,
-        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
-    ) -> anyhow::Result<Answer> {
+    ) -> Answer {
         self.role_state.accept(
             client_id,
             snownet::Offer {
@@ -60,11 +54,6 @@ impl GatewayTunnel {
                 },
             },
             client,
-            ipv4,
-            ipv6,
-            domain,
-            expires_at,
-            resource,
             Instant::now(),
         )
     }
@@ -75,13 +64,22 @@ impl GatewayTunnel {
 
     pub fn allow_access(
         &mut self,
-        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
         client: ClientId,
-        expires_at: Option<DateTime<Utc>>,
+        ipv4: Ipv4Addr,
+        ipv6: Ipv6Addr,
         domain: Option<(DomainName, Vec<IpAddr>)>,
+        expires_at: Option<DateTime<Utc>>,
+        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
     ) -> anyhow::Result<()> {
-        self.role_state
-            .allow_access(resource, client, expires_at, domain, Instant::now())
+        self.role_state.allow_access(
+            client,
+            ipv4,
+            ipv6,
+            domain,
+            expires_at,
+            resource,
+            Instant::now(),
+        )
     }
 
     pub fn refresh_translation(
@@ -230,54 +228,19 @@ impl GatewayState {
     }
 
     /// Accept a connection request from a client.
-    #[allow(clippy::too_many_arguments)]
     pub fn accept(
         &mut self,
         client_id: ClientId,
         offer: snownet::Offer,
         client: PublicKey,
-        ipv4: Ipv4Addr,
-        ipv6: Ipv6Addr,
-        domain: Option<(DomainName, Vec<IpAddr>)>,
-        expires_at: Option<DateTime<Utc>>,
-        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
         now: Instant,
-    ) -> anyhow::Result<Answer> {
-        match (&domain, &resource) {
-            (Some((domain, _)), ResourceDescription::Dns(r)) => {
-                if !crate::dns::is_subdomain(domain, &r.domain) {
-                    bail!(
-                        "Requested domain '{domain}' isn't a sub-domain of resource address '{}'",
-                        r.domain
-                    );
-                }
-            }
-            (None, ResourceDescription::Dns(_)) => {
-                bail!("Cannot setup connection for DNS resource without domain")
-            }
-            _ => {}
-        }
-
+    ) -> Answer {
         let answer = self.node.accept_connection(client_id, offer, client, now);
 
-        let mut peer = ClientOnGateway::new(client_id, ipv4, ipv6);
-
-        peer.add_resource(
-            resource.addresses(),
-            resource.id(),
-            resource.filters(),
-            expires_at,
-            domain.clone().map(|(n, _)| n),
-        );
-
-        peer.assign_proxies(&resource, domain, now)?;
-
-        self.peers.insert(peer, &[ipv4.into(), ipv6.into()]);
-
-        Ok(Answer {
+        Answer {
             username: answer.credentials.username,
             password: answer.credentials.password,
-        })
+        }
     }
 
     pub fn refresh_translation(
@@ -295,12 +258,15 @@ impl GatewayState {
         peer.refresh_translation(name, resource_id, resolved_ips, now);
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn allow_access(
         &mut self,
-        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
         client: ClientId,
-        expires_at: Option<DateTime<Utc>>,
+        ipv4: Ipv4Addr,
+        ipv6: Ipv6Addr,
         domain: Option<(DomainName, Vec<IpAddr>)>,
+        expires_at: Option<DateTime<Utc>>,
+        resource: ResourceDescription<ResolvedResourceDescriptionDns>,
         now: Instant,
     ) -> anyhow::Result<()> {
         match (&domain, &resource) {
@@ -318,7 +284,10 @@ impl GatewayState {
             _ => {}
         }
 
-        let peer = self.peers.get_mut(&client).context("Unknown client")?;
+        let peer = self
+            .peers
+            .entry(client)
+            .or_insert_with(|| ClientOnGateway::new(client, ipv4, ipv6));
 
         peer.assign_proxies(&resource, domain.clone(), now)?;
         peer.add_resource(
@@ -328,6 +297,8 @@ impl GatewayState {
             expires_at,
             domain.map(|(n, _)| n),
         );
+        self.peers.add_ip(&client, &ipv4.into());
+        self.peers.add_ip(&client, &ipv6.into());
 
         tracing::info!(%client, resource = %resource.id(), expires = ?expires_at.map(|e| e.to_rfc3339()), "Allowing access to resource");
         Ok(())

rust/connlib/tunnel/src/peer_store.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{hash_map::Entry, HashMap};
 use std::hash::Hash;
 use std::net::IpAddr;
 
@@ -67,6 +67,10 @@ where
         old_peer
     }
 
+    pub(crate) fn entry(&mut self, id: TId) -> Entry<'_, TId, P> {
+        self.peer_by_id.entry(id)
+    }
+
     pub(crate) fn remove(&mut self, id: &TId) -> Option<P> {
         self.id_by_ip.retain(|_, r_id| r_id != id);
         self.peer_by_id.remove(id)

rust/connlib/tunnel/src/tests/sut.rs

@@ -662,10 +662,12 @@ impl TunnelTest {
                 gateway
                     .exec_mut(|g| {
                         g.sut.allow_access(
-                            resource,
                             self.client.inner().id,
-                            None,
+                            self.client.inner().sut.tunnel_ip4().unwrap(),
+                            self.client.inner().sut.tunnel_ip6().unwrap(),
                             maybe_domain.map(|r| (r.name, r.proxy_ips)),
+                            None,
+                            resource,
                             now,
                         )
                     })
@@ -702,10 +704,12 @@ impl TunnelTest {
                     return;
                 };
 
+                let client_id = self.client.inner().id;
+
                 let answer = gateway
                     .exec_mut(|g| {
-                        g.sut.accept(
-                            self.client.inner().id,
+                        let answer = g.sut.accept(
+                            client_id,
                             snownet::Offer {
                                 session_key: preshared_key.expose_secret().0.into(),
                                 credentials: snownet::Credentials {
@@ -714,6 +718,10 @@ impl TunnelTest {
                                 },
                             },
                             self.client.inner().sut.public_key(),
+                            now,
+                        );
+                        g.sut.allow_access(
+                            client_id,
                             self.client.inner().sut.tunnel_ip4().unwrap(),
                             self.client.inner().sut.tunnel_ip6().unwrap(),
                             maybe_domain
@@ -722,7 +730,9 @@ impl TunnelTest {
                             None, // TODO: How to generate expiry?
                             resource,
                             now,
-                        )
+                        )?;
+
+                        anyhow::Ok(answer)
                     })
                     .unwrap();
 

rust/gateway/src/eventloop.rs

@@ -265,24 +265,28 @@ impl Eventloop {
             .inspect_err(|e| tracing::debug!(client = %req.client.id, reference = %req.reference, "DNS resolution timed out as part of connection request: {e}"))
             .unwrap_or_default();
 
-        match self.tunnel.accept(
+        let answer = self.tunnel.accept(
             req.client.id,
             req.client.peer.preshared_key,
             req.client.payload.ice_parameters,
             PublicKey::from(req.client.peer.public_key.0),
+        );
+
+        match self.tunnel.allow_access(
+            req.client.id,
             req.client.peer.ipv4,
             req.client.peer.ipv6,
             req.client.payload.domain.as_ref().map(|r| r.as_tuple()),
             req.expires_at,
             req.resource.into_resolved(addresses.clone()),
         ) {
-            Ok(accepted) => {
+            Ok(()) => {
                 self.portal.send(
                     PHOENIX_TOPIC,
                     EgressMessages::ConnectionReady(ConnectionReady {
                         reference: req.reference,
                         gateway_payload: GatewayResponse::ConnectionAccepted(ConnectionAccepted {
-                            ice_parameters: accepted,
+                            ice_parameters: answer,
                             domain_response: req.client.payload.domain.map(|r| {
                                 connlib_shared::messages::DomainResponse {
                                     domain: r.name(),
@@ -311,10 +315,12 @@ impl Eventloop {
 
         if let (Ok(()), Some(resolve_request)) = (
             self.tunnel.allow_access(
-                req.resource.into_resolved(addresses.clone()),
                 req.client_id,
-                req.expires_at,
+                req.client_ipv4,
+                req.client_ipv6,
                 req.payload.as_ref().map(|r| r.as_tuple()),
+                req.expires_at,
+                req.resource.into_resolved(addresses.clone()),
             ),
             req.payload,
         ) {

rust/gateway/src/messages.rs

@@ -1,4 +1,7 @@
-use std::{collections::BTreeSet, net::IpAddr};
+use std::{
+    collections::BTreeSet,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+};
 
 use chrono::{serde::ts_seconds_option, DateTime, Utc};
 use connlib_shared::{
@@ -83,7 +86,7 @@ impl ResolveRequest {
     }
 }
 
-#[derive(Debug, Deserialize, Clone, PartialEq, Eq)]
+#[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct AllowAccess {
     pub client_id: ClientId,
     pub resource: ResourceDescription,
@@ -92,6 +95,10 @@ pub struct AllowAccess {
     pub payload: Option<ResolveRequest>,
     #[serde(rename = "ref")]
     pub reference: String,
+    /// Tunnel IPv4 address.
+    pub client_ipv4: Ipv4Addr,
+    /// Tunnel IPv6 address.
+    pub client_ipv6: Ipv6Addr,
 }
 
 #[derive(Debug, Deserialize, Clone, PartialEq, Eq)]

scripts/tests/direct-dns-two-resources.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+
+# The integration tests call this to test Linux DNS control, using the `/etc/resolv.conf`
+# method which only works well inside Alpine Docker containers.
+
+source "./scripts/tests/lib.sh"
+
+RESOURCE1=dns.httpbin
+RESOURCE2=download.httpbin
+
+echo "# Try to ping httpbin as DNS resource 1"
+client_ping_resource "$RESOURCE1"
+
+echo "# Try to ping httpbin as DNS resource 2"
+client_ping_resource "$RESOURCE2"