fix(connlib): handle resources changing site (#10604)

Similar to how resources can be edited to change their address, IP stack
or other properties, they can also be moved between different sites.
Currently, `connlib` requires the portal to explicitly remove the
resource and then re-add it for this to work.

Our system gets more robust if we also detect that the sites of a
resource have changed and handle it like other addressability changes.

To ensure that this works correctly, we also extend the proptests to
simulate addressability changes of resources.

Resolves: #9881
Related: #10593

.github/workflows/_rust.yml

@@ -115,6 +115,7 @@ jobs:
           rg --count --no-ignore "Forwarding query for DNS resource to corresponding site" "$TESTCASES_DIR"
           rg --count --no-ignore "Revoking resource authorization" "$TESTCASES_DIR"
           rg --count --no-ignore "Re-seeding records for DNS resources" "$TESTCASES_DIR"
+          rg --count --no-ignore "Resource is known but its addressability changed" "$TESTCASES_DIR"
 
           # Make sure we are recovering from ICE disconnect
           rg --count --no-ignore "State change \(got new possible\): Disconnected -> Checking" "$TESTCASES_DIR"

rust/connlib/tunnel/proptest-regressions/tests.txt

@@ -214,3 +214,8 @@ cc 5f9aab8caf96d39aae4869a153a16d32193ad91a9c0095233e3a145b85780e08
 cc 6c0a337606eb3df1ecee7aec3adc5dee37860c822d889f40e9f4b078c101e628
 cc 8a184391bcfcc341c5e12149880bfae1850ecf6d49195e813ef3f1ac7ca20cba
 cc 3460ab82122404d6815b447f0c8e403a00f35fa7d493e5f921583936ef51ac5e
+cc 3088baa8843b8a7531a4b4cb164d769c1725c77d2b99ff0d7b81acae764c24ec
+cc 3467fb0a9697b7b1221b46558d998b3689bdce49944de7fcdc2627e1fbbc3771
+cc 3bdd819cda2577278b0372cb7598418227ecab83271c48f5b28dc192f766061e
+cc 764c22e664da06820cd02cba259196edeec94cce45e450959ce9354be7bc9f1c
+cc 04193ee1047f542c469aa0893bf636df9c317943022d922e231de3e821b39486

rust/connlib/tunnel/src/client.rs

@@ -1,9 +1,11 @@
 mod dns_cache;
 mod dns_resource_nat;
 mod gateway_on_client;
+mod pending_tun_update;
 mod resource;
 
 pub(crate) use crate::client::gateway_on_client::GatewayOnClient;
+use crate::client::pending_tun_update::PendingTunUpdate;
 #[cfg(all(feature = "proptest", test))]
 pub(crate) use resource::DnsResource;
 pub(crate) use resource::{CidrResource, InternetResource, Resource};
@@ -144,6 +146,7 @@ pub struct ClientState {
     /// We use this as a hint to the portal to re-connect us to the same gateway for a resource.
     recently_connected_gateways: LruCache<GatewayId, ()>,
 
+    pending_tun_update: Option<PendingTunUpdate>,
     buffered_events: VecDeque<ClientEvent>,
     buffered_packets: VecDeque<IpPacket>,
     buffered_transmits: VecDeque<Transmit>,
@@ -224,6 +227,7 @@ impl ClientState {
             tcp_dns_streams_by_upstream_and_query_id: Default::default(),
             pending_flows: Default::default(),
             dns_resource_nat: Default::default(),
+            pending_tun_update: Default::default(),
         }
     }
 
@@ -1557,13 +1561,16 @@ impl ClientState {
         self.stub_resolver
             .set_search_domain(new_tun_config.search_domain.clone());
 
-        // Ensure we don't emit multiple interface updates in a row.
-        self.buffered_events
-            .retain(|e| !matches!(e, ClientEvent::TunInterfaceUpdated(_)));
-
-        self.tun_config = Some(new_tun_config.clone());
-        self.buffered_events
-            .push_back(ClientEvent::TunInterfaceUpdated(new_tun_config));
+        match self.pending_tun_update.as_mut() {
+            Some(pending) => pending.update_want(new_tun_config.clone()),
+            None => {
+                self.pending_tun_update = Some(PendingTunUpdate::new(
+                    self.tun_config.clone(),
+                    new_tun_config.clone(),
+                ))
+            }
+        }
+        self.tun_config = Some(new_tun_config);
 
         self.initialise_tcp_dns_client(); // We must reset the TCP DNS client because changed CIDR resources (and thus changed routes) might affect which site we connect to.
         self.initialise_tcp_dns_server();
@@ -1639,6 +1646,12 @@ impl ClientState {
     }
 
     pub(crate) fn poll_event(&mut self) -> Option<ClientEvent> {
+        if let Some(pending_tun_update) = self.pending_tun_update.take()
+            && let Some(event) = pending_tun_update.into_event()
+        {
+            return Some(event);
+        }
+
         self.buffered_events
             .pop_front()
             .or_else(|| match self.stub_resolver.poll_event()? {
@@ -1732,9 +1745,12 @@ impl ClientState {
 
         if let Some(resource) = self.resources_by_id.get(&new_resource.id()) {
             let resource_addressability_changed = resource.has_different_address(&new_resource)
-                || resource.has_different_ip_stack(&new_resource);
+                || resource.has_different_ip_stack(&new_resource)
+                || resource.has_different_site(&new_resource);
 
             if resource_addressability_changed {
+                tracing::debug!(rid = %new_resource.id(), "Resource is known but its addressability changed");
+
                 self.remove_resource(resource.id(), now);
             }
         }

rust/connlib/tunnel/src/client/pending_tun_update.rs

@@ -0,0 +1,32 @@
+use crate::{ClientEvent, TunConfig};
+
+/// Acts as a "buffer" for updates to the TUN interface that need to be applied.
+///
+/// When adding one or more resources, it can happen that multiple updates to the TUN device need to be issued.
+/// In order to not actually send multiple updates, we buffer them in here.
+///
+/// In the event that the very last one ends up being the state we are already in, no update is issued at all.
+pub struct PendingTunUpdate {
+    current: Option<TunConfig>,
+    want: TunConfig,
+}
+
+impl PendingTunUpdate {
+    pub fn new(current: Option<TunConfig>, want: TunConfig) -> Self {
+        Self { current, want }
+    }
+
+    pub fn update_want(&mut self, want: TunConfig) {
+        self.want = want;
+    }
+
+    pub fn into_event(self) -> Option<ClientEvent> {
+        if let Some(current) = self.current
+            && current == self.want
+        {
+            return None;
+        };
+
+        Some(ClientEvent::TunInterfaceUpdated(self.want))
+    }
+}

rust/connlib/tunnel/src/client/resource.rs

@@ -176,6 +176,10 @@ impl Resource {
         }
     }
 
+    pub fn has_different_site(&self, other: &Resource) -> bool {
+        self.sites() != other.sites()
+    }
+
     pub fn addresses(&self) -> Vec<IpNetwork> {
         match self {
             Resource::Dns(_) => vec![],
@@ -194,6 +198,24 @@ impl Resource {
             Resource::Internet(r) => ResourceView::Internet(r.with_status(status)),
         }
     }
+
+    #[cfg(all(test, feature = "proptest"))]
+    pub fn with_new_site(self, site: Site) -> Self {
+        match self {
+            Resource::Dns(r) => Self::Dns(DnsResource {
+                sites: vec![site],
+                ..r
+            }),
+            Resource::Cidr(r) => Self::Cidr(CidrResource {
+                sites: vec![site],
+                ..r
+            }),
+            Resource::Internet(r) => Self::Internet(InternetResource {
+                sites: vec![site],
+                ..r
+            }),
+        }
+    }
 }
 
 impl TryFrom<ResourceDescription> for Resource {

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

@@ -6,7 +6,7 @@ use super::{
 };
 use crate::client;
 use crate::{dns::is_subdomain, proptest::relay_id};
-use connlib_model::{GatewayId, RelayId, StaticSecret};
+use connlib_model::{GatewayId, RelayId, Site, StaticSecret};
 use dns_types::{DomainName, RecordType};
 use ip_network::{Ipv4Network, Ipv6Network};
 use itertools::Itertools;
@@ -230,6 +230,29 @@ impl ReferenceState {
                 state.all_resources_not_known_to_client(),
                 |resource_ids| sample::select(resource_ids).prop_map(Transition::AddResource),
             )
+            .with_if_not_empty(1, state.cidr_resources_on_client(), |resources| {
+                (sample::select(resources), cidr_resource_address()).prop_map(
+                    |(resource, new_address)| Transition::ChangeCidrResourceAddress {
+                        resource,
+                        new_address,
+                    },
+                )
+            })
+            .with_if_not_empty(
+                1,
+                (
+                    state.cidr_and_dns_resources_on_client(),
+                    state.regular_sites(),
+                ),
+                |(resources, sites)| {
+                    (sample::select(resources), sample::select(sites)).prop_map(
+                        |(resource, new_site)| Transition::MoveResourceToNewSite {
+                            resource,
+                            new_site,
+                        },
+                    )
+                },
+            )
             .with_if_not_empty(1, state.client.inner().all_resource_ids(), |resource_ids| {
                 sample::select(resource_ids).prop_map(Transition::RemoveResource)
             })
@@ -406,6 +429,36 @@ impl ReferenceState {
                     client.remove_resource(id);
                 });
             }
+            Transition::ChangeCidrResourceAddress {
+                resource,
+                new_address,
+            } => {
+                state
+                    .portal
+                    .change_address_of_cidr_resource(resource.id, *new_address);
+
+                let new_resource = client::CidrResource {
+                    address: *new_address,
+                    ..resource.clone()
+                };
+
+                state.client.exec_mut(|c| c.add_cidr_resource(new_resource));
+            }
+            Transition::MoveResourceToNewSite { resource, new_site } => {
+                state
+                    .portal
+                    .move_resource_to_new_site(resource.id(), new_site.clone());
+
+                state
+                    .client
+                    .exec_mut(|c| match resource.clone().with_new_site(new_site.clone()) {
+                        client::Resource::Dns(r) => c.add_dns_resource(r),
+                        client::Resource::Cidr(r) => c.add_cidr_resource(r),
+                        client::Resource::Internet(_) => {
+                            tracing::error!("Internet Resource cannot move site");
+                        }
+                    })
+            }
             Transition::SetInternetResourceState(active) => state.client.exec_mut(|client| {
                 client.set_internet_resource_state(*active);
             }),
@@ -526,6 +579,14 @@ impl ReferenceState {
 
                 true
             }
+            Transition::ChangeCidrResourceAddress {
+                resource,
+                new_address,
+            } => resource.address != *new_address && state.client.inner().has_resource(resource.id),
+            Transition::MoveResourceToNewSite { resource, new_site } => {
+                resource.sites() != BTreeSet::from([new_site])
+                    && state.client.inner().has_resource(resource.id())
+            }
             Transition::SetInternetResourceState(_) => true,
             Transition::SendIcmpPacket {
                 src,
@@ -790,6 +851,40 @@ impl ReferenceState {
         all_resources
     }
 
+    fn cidr_and_dns_resources_on_client(&self) -> Vec<client::Resource> {
+        let mut all_resources = self.portal.all_resources();
+        all_resources.retain(|r| {
+            matches!(r, client::Resource::Cidr(_) | client::Resource::Dns(_))
+                && self.client.inner().has_resource(r.id())
+        });
+
+        all_resources
+    }
+
+    fn cidr_resources_on_client(&self) -> Vec<client::CidrResource> {
+        self.portal
+            .all_resources()
+            .into_iter()
+            .flat_map(|r| match r {
+                client::Resource::Cidr(r) => Some(r),
+                client::Resource::Dns(_) | client::Resource::Internet(_) => None,
+            })
+            .filter(|r| self.client.inner().has_resource(r.id))
+            .collect()
+    }
+
+    fn regular_sites(&self) -> Vec<Site> {
+        let all_sites = self
+            .portal
+            .all_resources()
+            .into_iter()
+            .filter(|r| !matches!(r, client::Resource::Internet(_)))
+            .flat_map(|r| r.sites().into_iter().cloned().collect::<Vec<_>>())
+            .collect::<BTreeSet<_>>();
+
+        Vec::from_iter(all_sites)
+    }
+
     fn connected_gateway_ipv4_ips(&self) -> Vec<Ipv4Network> {
         self.client
             .inner()

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

@@ -648,21 +648,46 @@ impl RefClient {
     }
 
     pub(crate) fn add_cidr_resource(&mut self, r: CidrResource) {
-        self.resources.push(Resource::Cidr(r.clone()));
+        let address = r.address;
+        let r = Resource::Cidr(r);
+
+        if let Some(existing) = self
+            .resources
+            .iter()
+            .find(|existing| existing.id() == r.id())
+            && (existing.has_different_address(&r) || existing.has_different_site(&r))
+        {
+            self.remove_resource(&existing.id());
+        }
+
+        self.resources.push(r.clone());
         self.cidr_resources = self.recalculate_cidr_routes();
 
-        match r.address {
+        match address {
             IpNetwork::V4(v4) => {
-                self.ipv4_routes.insert(r.id, v4);
+                self.ipv4_routes.insert(r.id(), v4);
             }
             IpNetwork::V6(v6) => {
-                self.ipv6_routes.insert(r.id, v6);
+                self.ipv6_routes.insert(r.id(), v6);
             }
         }
     }
 
     pub(crate) fn add_dns_resource(&mut self, r: DnsResource) {
-        self.resources.push(Resource::Dns(r));
+        let r = Resource::Dns(r);
+
+        if let Some(existing) = self
+            .resources
+            .iter()
+            .find(|existing| existing.id() == r.id())
+            && (existing.has_different_address(&r)
+                || existing.has_different_ip_stack(&r)
+                || existing.has_different_site(&r))
+        {
+            self.remove_resource(&existing.id());
+        }
+
+        self.resources.push(r);
     }
 
     /// Re-adds all resources in the order they have been initially added.
@@ -844,7 +869,7 @@ impl RefClient {
         let previous = self.site_status.insert(site.id, ResourceStatus::Online);
 
         if previous.is_none_or(|s| s != ResourceStatus::Online) {
-            tracing::debug!(%rid, "Resource is now online");
+            tracing::debug!(%rid, sid = %site.id, "Resource is now online");
         }
     }
 

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

@@ -9,7 +9,7 @@ use crate::messages::DnsServer;
 use crate::{IPV4_TUNNEL, IPV6_TUNNEL, proptest::*};
 use connlib_model::{RelayId, Site};
 use dns_types::{DomainName, OwnedRecordData};
-use ip_network::{Ipv4Network, Ipv6Network};
+use ip_network::{IpNetwork, Ipv4Network, Ipv6Network};
 use itertools::Itertools;
 use prop::sample;
 use proptest::{collection, prelude::*};
@@ -267,10 +267,14 @@ fn cidr_resource_outside_reserved_ranges(
     sites: impl Strategy<Value = Site>,
 ) -> impl Strategy<Value = CidrResource> {
     cidr_resource(
-        non_reserved_ip().prop_flat_map(move |ip| ip_network(ip, 8)), sites.prop_map(|s| vec![s]))
+        cidr_resource_address(), sites.prop_map(|s| vec![s]))
         .prop_filter("resource must not be in the documentation range because we use those for host addresses and DNS IPs", |r| !r.address.is_documentation())
 }
 
+pub(crate) fn cidr_resource_address() -> impl Strategy<Value = IpNetwork> {
+    non_reserved_ip().prop_flat_map(move |ip| ip_network(ip, 8))
+}
+
 fn internet_resource(site: impl Strategy<Value = Site>) -> impl Strategy<Value = InternetResource> {
     crate::proptest::internet_resource(site.prop_map(|s| vec![s]))
 }

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

@@ -10,9 +10,10 @@ use crate::{
     client::DnsResource,
     messages::{DnsServer, gateway},
 };
-use connlib_model::GatewayId;
+use connlib_model::{GatewayId, Site};
 use connlib_model::{ResourceId, SiteId};
 use dns_types::DomainName;
+use ip_network::IpNetwork;
 use itertools::Itertools;
 use proptest::{
     collection,
@@ -227,6 +228,37 @@ impl StubPortal {
             .map(|(gid, _, _)| *gid)
     }
 
+    pub(crate) fn change_address_of_cidr_resource(
+        &mut self,
+        rid: ResourceId,
+        new_address: IpNetwork,
+    ) {
+        if let Some(resource) = self.cidr_resources.get_mut(&rid) {
+            resource.address = new_address;
+            return;
+        }
+
+        tracing::error!(%rid, "Unknown resource");
+    }
+
+    pub(crate) fn move_resource_to_new_site(&mut self, rid: ResourceId, site: Site) {
+        if let Some(resource) = self.cidr_resources.get_mut(&rid) {
+            self.sites_by_resource.insert(rid, site.id);
+            resource.sites = vec![site];
+            return;
+        }
+
+        if let Some(resource) = self.dns_resources.get_mut(&rid) {
+            self.sites_by_resource.insert(rid, site.id);
+            resource.sites = vec![site];
+            return;
+        }
+
+        if self.internet_resource.id == rid {
+            tracing::error!("Internet Resource cannot change site");
+        }
+    }
+
     pub(crate) fn gateways(
         &self,
     ) -> impl Strategy<Value = BTreeMap<GatewayId, Host<RefGateway>>> + use<> {

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

@@ -8,7 +8,7 @@ use super::sim_net::{Host, HostId, RoutingTable};
 use super::sim_relay::SimRelay;
 use super::stub_portal::StubPortal;
 use super::transition::{Destination, DnsQuery};
-use crate::client::Resource;
+use crate::client;
 use crate::dns::is_subdomain;
 use crate::messages::{IceCredentials, Key, SecretKey};
 use crate::tests::assertions::*;
@@ -130,7 +130,7 @@ impl TunnelTest {
                 state.client.exec_mut(|c| {
                     // Flush DNS.
                     match &resource {
-                        Resource::Dns(r) => {
+                        client::Resource::Dns(r) => {
                             c.dns_records.retain(|domain, _| {
                                 if is_subdomain(domain, &r.address) {
                                     return false;
@@ -139,13 +139,43 @@ impl TunnelTest {
                                 true
                             });
                         }
-                        Resource::Cidr(_) => {}
-                        Resource::Internet(_) => {}
+                        client::Resource::Cidr(_) => {}
+                        client::Resource::Internet(_) => {}
                     }
 
                     c.sut.add_resource(resource, now);
                 });
             }
+            Transition::ChangeCidrResourceAddress {
+                resource,
+                new_address,
+            } => {
+                let new_resource = client::Resource::Cidr(client::CidrResource {
+                    address: new_address,
+                    ..resource.clone()
+                });
+
+                if let Some(gateway) = ref_state
+                    .portal
+                    .gateway_for_resource(new_resource.id())
+                    .and_then(|gid| state.gateways.get_mut(gid))
+                {
+                    gateway.exec_mut(|g| {
+                        g.sut
+                            .remove_access(&state.client.inner().id, &new_resource.id(), now)
+                    })
+                }
+                state
+                    .client
+                    .exec_mut(|c| c.sut.add_resource(new_resource, now));
+            }
+            Transition::MoveResourceToNewSite { resource, new_site } => {
+                let new_resource = resource.with_new_site(new_site);
+
+                state
+                    .client
+                    .exec_mut(|c| c.sut.add_resource(new_resource, now));
+            }
             Transition::RemoveResource(rid) => {
                 state.client.exec_mut(|c| c.sut.remove_resource(rid, now));
 

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

@@ -1,9 +1,10 @@
 use crate::{
-    client::{IPV4_RESOURCES, IPV6_RESOURCES, Resource},
+    client::{CidrResource, IPV4_RESOURCES, IPV6_RESOURCES, Resource},
     proptest::{host_v4, host_v6},
 };
-use connlib_model::{RelayId, ResourceId};
+use connlib_model::{RelayId, ResourceId, Site};
 use dns_types::{DomainName, RecordType};
+use ip_network::IpNetwork;
 
 use super::{
     reference::PrivateKey,
@@ -25,6 +26,13 @@ pub(crate) enum Transition {
     AddResource(Resource),
     /// Remove a resource on the client.
     RemoveResource(ResourceId),
+    /// Change the address of a CIDR resource.
+    ChangeCidrResourceAddress {
+        resource: CidrResource,
+        new_address: IpNetwork,
+    },
+    /// Move a CIDR/DNS resource to a new site.
+    MoveResourceToNewSite { resource: Resource, new_site: Site },
 
     /// Toggle the Internet Resource on / off
     SetInternetResourceState(bool),