From e6b0cf842a79daa17152071d9e8ad66a9b1c3915 Mon Sep 17 00:00:00 2001
From: Cecylia Bocovich <cohosh@torproject.org>
Date: Sat, 9 Mar 2024 19:14:14 -0500
Subject: [PATCH] Pass available_bridge by value rather than reference

Allows us to pop() from extra_bridges, which is less prone to errors than
calling both last() and then remove().
---
 crates/lox-distributor/src/lox_context.rs | 19 +++++++------------
 crates/lox-library/src/lib.rs             |  6 +++---
 crates/lox-library/src/tests.rs           | 12 ++++++------
 3 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/crates/lox-distributor/src/lox_context.rs b/crates/lox-distributor/src/lox_context.rs
index 98159e3..151a3ce 100644
--- a/crates/lox-distributor/src/lox_context.rs
+++ b/crates/lox-distributor/src/lox_context.rs
@@ -251,12 +251,6 @@ impl LoxServerContext {
         return_bridges
     }
 
-    pub fn remove_single_bridge(&self) {
-        let mut extra_bridges = self.extra_bridges.lock().unwrap();
-        let length = extra_bridges.len();
-        _ = extra_bridges.remove(length - 1)
-    }
-
     // Add extra_bridges to the Lox bridge table as open invitation bridges
     // TODO: Add some consideration for whether or not bridges should be sorted as
     // open invitation buckets or hot spare buckets
@@ -301,13 +295,14 @@ impl LoxServerContext {
     // available bridges or from a spare bucket
     pub fn replace_with_new(&self, bridgeline: BridgeLine) -> lox_library::ReplaceSuccess {
         let mut ba_obj = self.ba.lock().unwrap();
-        let eb_obj = self.extra_bridges.lock().unwrap();
-        let available_bridge = eb_obj.last();
+        let mut eb_obj = self.extra_bridges.lock().unwrap();
+        let available_bridge = eb_obj.pop();
         let result = ba_obj.bridge_replace(&bridgeline, available_bridge);
-        // .last() doesn't actually remove the object so we still have to do that if the bridge was
-        // replaced with an available bridge
-        if result == lox_library::ReplaceSuccess::Replaced && eb_obj.len() > 0 {
-            self.remove_single_bridge();
+        if result != lox_library::ReplaceSuccess::Replaced {
+            if let Some(bridge) = available_bridge {
+                // If available bridge wasn't removed, return it
+                eb_obj.push(bridge);
+            }
         }
         result
     }
diff --git a/crates/lox-library/src/lib.rs b/crates/lox-library/src/lib.rs
index b2b73a1..6152d97 100644
--- a/crates/lox-library/src/lib.rs
+++ b/crates/lox-library/src/lib.rs
@@ -540,7 +540,7 @@ impl BridgeAuth {
     pub fn bridge_replace(
         &mut self,
         bridge: &BridgeLine,
-        available_bridge: Option<&BridgeLine>,
+        available_bridge: Option<BridgeLine>,
     ) -> ReplaceSuccess {
         let mut res = ReplaceSuccess::NotFound;
         let reachable_bridges = &self.bridge_table.reachable.clone();
@@ -554,12 +554,12 @@ impl BridgeAuth {
                     None => return ReplaceSuccess::NotFound,
                 };
                 assert!(bridgelines[*offset] == *bridge);
-                bridgelines[*offset] = *replacement;
+                bridgelines[*offset] = replacement;
                 self.bridge_table.buckets.insert(*bucketnum, bridgelines);
                 // Remove the bridge from the reachable bridges and add new bridge
                 self.bridge_table
                     .reachable
-                    .insert(*replacement, positions.clone());
+                    .insert(replacement, positions.clone());
                 // Remove the bridge from the bucket
                 self.bridge_table.reachable.remove(bridge);
             }
diff --git a/crates/lox-library/src/tests.rs b/crates/lox-library/src/tests.rs
index f74284d..1e7b085 100644
--- a/crates/lox-library/src/tests.rs
+++ b/crates/lox-library/src/tests.rs
@@ -1063,21 +1063,21 @@ fn test_bridge_replace() {
         match case {
             "not found" => {
                 // Case zero: bridge to be replaced is not in the bridgetable
-                let random_bridgeline = &BridgeLine::random();
+                let random_bridgeline = BridgeLine::random();
                 assert!(
-                    !th.ba.bridge_table.reachable.contains_key(random_bridgeline),
+                    !th.ba.bridge_table.reachable.contains_key(&random_bridgeline),
                     "Random bridgeline happens to be in the bridge_table (and should not be)"
                 );
                 assert!(
                     th.ba
-                        .bridge_replace(random_bridgeline, Some(random_bridgeline))
+                        .bridge_replace(&random_bridgeline, Some(random_bridgeline))
                         == ReplaceSuccess::NotFound,
                     "Bridge should be marked as NotFound"
                 );
             }
             "available" => {
                 // Case one: available_bridge != null
-                let random_bridgeline = &BridgeLine::random();
+                let random_bridgeline = BridgeLine::random();
                 let unallocated_bridgeline = &BridgeLine::random();
                 th.ba
                     .bridge_table
@@ -1087,7 +1087,7 @@ fn test_bridge_replace() {
                     th.ba
                         .bridge_table
                         .reachable
-                        .get(random_bridgeline)
+                        .get(&random_bridgeline)
                         .is_none(),
                     "Random bridge already in table"
                 );
@@ -1101,7 +1101,7 @@ fn test_bridge_replace() {
                     th.ba
                         .bridge_table
                         .reachable
-                        .get(random_bridgeline)
+                        .get(&random_bridgeline)
                         .is_some(),
                     "Replacement bridge not added to reachable bridges"
                 );