From 10c9f622f3a5a92a6e7fdad933bf2063138ff63b Mon Sep 17 00:00:00 2001
From: Gabi <gabrielalejandro7@gmail.com>
Date: Thu, 15 Feb 2024 22:23:44 -0300
Subject: [PATCH] fix(snownet): generate candidates only after we accept the
 ICE answer (#3658)

This is done to delay the candidate generation after the gateway has
already received the request.

Since we already know the candidates in most cases, an optimization in
the future to reduce the number of round-trips to the gateway we can add
the candidates to the request connection message.
---
 rust/connlib/snownet/src/node.rs  | 14 +++++------
 rust/connlib/snownet/tests/lib.rs | 42 ++++++++++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs
index d3effd9ae..7c999ea6a 100644
--- a/rust/connlib/snownet/src/node.rs
+++ b/rust/connlib/snownet/src/node.rs
@@ -666,13 +666,6 @@ where
         let mut agent = IceAgent::new();
         agent.set_controlling(true);
 
-        self.seed_agent_with_local_candidates(
-            id,
-            &mut agent,
-            &allowed_stun_servers,
-            &allowed_turn_servers,
-        );
-
         let session_key = Secret::new(random());
         let ice_creds = agent.local_credentials();
 
@@ -714,6 +707,13 @@ where
             pass: answer.credentials.password,
         });
 
+        self.seed_agent_with_local_candidates(
+            id,
+            &mut agent,
+            &initial.stun_servers,
+            &initial.turn_servers,
+        );
+
         let connection = self.init_connection(
             agent,
             remote,
diff --git a/rust/connlib/snownet/tests/lib.rs b/rust/connlib/snownet/tests/lib.rs
index 81a2a8700..b5aad89a3 100644
--- a/rust/connlib/snownet/tests/lib.rs
+++ b/rust/connlib/snownet/tests/lib.rs
@@ -2,9 +2,11 @@ use boringtun::x25519::StaticSecret;
 use snownet::{ClientNode, Event, ServerNode};
 use std::{
     collections::HashSet,
-    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    iter,
+    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
     time::{Duration, Instant},
 };
+use str0m::{net::Protocol, Candidate};
 
 #[test]
 fn connection_times_out_after_10_seconds() {
@@ -36,6 +38,44 @@ fn answer_after_stale_connection_does_not_panic() {
     alice.accept_answer(1, bob.public_key(), answer);
 }
 
+#[test]
+fn only_generate_candidate_event_after_answer() {
+    let local_candidate = SocketAddr::new(IpAddr::from(Ipv4Addr::LOCALHOST), 10000);
+
+    let mut alice = ClientNode::<u64>::new(
+        StaticSecret::random_from_rng(rand::thread_rng()),
+        Instant::now(),
+    );
+
+    alice.add_local_host_candidate(local_candidate).unwrap();
+
+    let mut bob = ServerNode::<u64>::new(
+        StaticSecret::random_from_rng(rand::thread_rng()),
+        Instant::now(),
+    );
+
+    let offer = alice.new_connection(1, HashSet::new(), HashSet::new());
+
+    assert_eq!(
+        alice.poll_event(),
+        None,
+        "no event to be emitted before accepting the answer"
+    );
+
+    let answer =
+        bob.accept_connection(1, offer, alice.public_key(), HashSet::new(), HashSet::new());
+
+    alice.accept_answer(1, bob.public_key(), answer);
+
+    assert!(iter::from_fn(|| alice.poll_event()).any(|ev| ev
+        == Event::SignalIceCandidate {
+            connection: 1,
+            candidate: Candidate::host(local_candidate, Protocol::Udp)
+                .unwrap()
+                .to_sdp_string()
+        }));
+}
+
 #[test]
 fn second_connection_with_same_relay_reuses_allocation() {
     let mut alice = ClientNode::<u64>::new(