From 0e84ef8fee61722abfc7556086fec94ff77cea82 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 3 Sep 2024 13:45:04 -0700
Subject: [PATCH] test(connlib): track pending connections to gateways (#6497)

Instead of tracking pending connections to resources, we need to model
pending connections to gateways. The offending test seed has a CIDR
resource that is a DNS server and the Internet resources, both routed
via the same gateway.

When sending concurrent DNS queries to those resources, we need to track
which _gateways_ we are connecting to as a result to figure out which
queries get lost. In particular, only the _first_ resource to trigger a
connection to a gateway will be authorized. Subsequent queries will be
completely lost and require another packet to authorize the connection.

---------

Signed-off-by: Thomas Eizinger <thomas@eizinger.io>
Co-authored-by: Not Applicable <ReactorScram@users.noreply.github.com>
---
 rust/Cargo.lock                               |  4 +--
 .../tunnel/proptest-regressions/tests.txt     |  3 ++
 rust/connlib/tunnel/src/tests/reference.rs    | 33 ++++++++++++-------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index c747430dd..0115b7156 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -4656,7 +4656,7 @@ dependencies = [
 [[package]]
 name = "proptest"
 version = "1.5.0"
-source = "git+https://github.com/proptest-rs/proptest?branch=master#c012218897b41606a5f8a21718d44d10509cb502"
+source = "git+https://github.com/proptest-rs/proptest?branch=master#ff98819628a73005fe8181e7f5f7d6e1a89fd565"
 dependencies = [
  "bit-set",
  "bit-vec",
@@ -4675,7 +4675,7 @@ dependencies = [
 [[package]]
 name = "proptest-state-machine"
 version = "0.3.0"
-source = "git+https://github.com/proptest-rs/proptest?branch=master#c012218897b41606a5f8a21718d44d10509cb502"
+source = "git+https://github.com/proptest-rs/proptest?branch=master#ff98819628a73005fe8181e7f5f7d6e1a89fd565"
 dependencies = [
  "proptest",
 ]
diff --git a/rust/connlib/tunnel/proptest-regressions/tests.txt b/rust/connlib/tunnel/proptest-regressions/tests.txt
index 0127f2630..41767aa77 100644
--- a/rust/connlib/tunnel/proptest-regressions/tests.txt
+++ b/rust/connlib/tunnel/proptest-regressions/tests.txt
@@ -93,3 +93,6 @@ cc fa1ad96fcad83aa29156936f13a7722ccc1b349bc8c2022de4a6a20c3f4e9121
 cc 6aab25ffe3551557cec151eacac706eeb2f4cf7bbae55b0c14f52830e7a076ff
 cc a47613091d3c393fd278753a79795ccff97e3c60d7a20ed08d8abce7473ea643
 cc fa1ad96fcad83aa29156936f13a7722ccc1b349bc8c2022de4a6a20c3f4e9121
+cc 404b2d3e98f7911fd32b3b684b262ba01fffba8683a1068e4ddb6ed8cddcb076
+cc ffd2e02bfbab42e0afab972ec86889d775b7a6a3bf50ac0b93e3f0860af41b3a
+cc 160cd0b6d9070855886c9e5c53e330bfd7494dd37fdba3a54a964ef36d1e3619
diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs
index 16ee14eb4..19ab0c703 100644
--- a/rust/connlib/tunnel/src/tests/reference.rs
+++ b/rust/connlib/tunnel/src/tests/reference.rs
@@ -14,7 +14,7 @@ use domain::base::Rtype;
 use proptest::{prelude::*, sample};
 use proptest_state_machine::ReferenceStateMachine;
 use std::{
-    collections::{BTreeMap, BTreeSet, HashSet},
+    collections::{btree_map::Entry, BTreeMap, BTreeSet, HashSet},
     fmt, iter,
     net::{IpAddr, SocketAddr},
 };
@@ -320,7 +320,7 @@ impl ReferenceStateMachine for ReferenceState {
                 }
             }),
             Transition::SendDnsQueries(queries) => {
-                let mut pending_connections = HashSet::new();
+                let mut new_connections_via_gateways = BTreeMap::new();
 
                 for query in queries {
                     // Some queries get answered locally.
@@ -342,28 +342,37 @@ impl ReferenceStateMachine for ReferenceState {
                         continue;
                     };
 
-                    tracing::debug!("Expecting DNS query via resource");
-
-                    if pending_connections.contains(&resource) {
-                        // DNS server is a CIDR resource and a previous query of this batch is already triggering a connection.
-                        // That connection isn't ready yet so further queries to the same resource are dropped until then.
+                    let Some(gateway) = state.portal.gateway_for_resource(resource).copied() else {
+                        tracing::error!("Unknown gateway for resource");
                         continue;
-                    }
+                    };
+
+                    tracing::debug!(%resource, %gateway, "Expecting DNS query via resource");
 
                     if !state
                         .client
                         .inner()
                         .is_connected_to_internet_or_cidr(resource)
                     {
-                        state.client.exec_mut(|client| {
-                            client.connect_to_internet_or_cidr_resource(resource)
-                        });
-                        pending_connections.insert(resource);
+                        // As part of batch-processing DNS queries, only the first resource per gateway will be connected / authorized.
+                        match new_connections_via_gateways.entry(gateway) {
+                            Entry::Vacant(v) => {
+                                v.insert(resource);
+                            }
+                            Entry::Occupied(_) => {}
+                        };
+
                         continue;
                     }
 
                     state.client.exec_mut(|client| client.on_dns_query(query));
                 }
+
+                for (_, resource) in new_connections_via_gateways {
+                    state
+                        .client
+                        .exec_mut(|client| client.connect_to_internet_or_cidr_resource(resource));
+                }
             }
             Transition::SendICMPPacketToNonResourceIp {
                 src,