From b92cd242e7265b018b4709a660d8b535e1831e65 Mon Sep 17 00:00:00 2001
From: onyinyang <onyinyang@torproject.org>
Date: Wed, 5 Apr 2023 16:26:23 -0400
Subject: [PATCH] Ensure that left over bridges (not grouped into a bucket) are
 properly handled

---
 crates/lox-distributor/src/lox_context.rs | 24 +++++++++++++++-
 crates/lox-distributor/src/main.rs        | 35 +++++++++++++++++------
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/crates/lox-distributor/src/lox_context.rs b/crates/lox-distributor/src/lox_context.rs
index 212b754..802a5a5 100644
--- a/crates/lox-distributor/src/lox_context.rs
+++ b/crates/lox-distributor/src/lox_context.rs
@@ -7,7 +7,7 @@ use hyper::{
 
 use lox::{
     BridgeAuth, BridgeDb, OPENINV_LENGTH,
-    bridge_table::{BridgeLine, ENC_BUCKET_BYTES},
+    bridge_table::{BridgeLine, ENC_BUCKET_BYTES, MAX_BRIDGES_PER_BUCKET},
     proto::{
         blockage_migration, check_blockage, issue_invite, level_up, migration, open_invite,
         redeem_invite, trust_promotion,
@@ -79,15 +79,37 @@ pub fn random() -> BridgeLine {
 pub struct LoxServerContext {
     pub db: Arc<Mutex<BridgeDb>>,
     pub ba: Arc<Mutex<BridgeAuth>>,
+    pub extra_bridges: Arc<Mutex<Vec<BridgeLine>>>,
 }
 
 impl LoxServerContext {
+
+    pub fn append_extra_bridges(&self, bridge: BridgeLine){
+        let mut extra_bridges = self.extra_bridges.lock().unwrap();
+        extra_bridges.push(bridge);
+    }
+
+    pub fn remove_extra_bridges(&self) -> [BridgeLine; MAX_BRIDGES_PER_BUCKET] {
+        let mut extra_bridges = self.extra_bridges.lock().unwrap();
+        [
+            extra_bridges.remove(0),
+            extra_bridges.remove(1),
+            extra_bridges.remove(2),
+        ]
+
+    }
+
     pub fn add_openinv_bucket(&self, bucket: [BridgeLine; 3]) {
         let mut ba_obj = self.ba.lock().unwrap();
         let mut db_obj = self.db.lock().unwrap();
         ba_obj.add_openinv_bridges(bucket, &mut db_obj);
     }
 
+    pub fn add_spare_bucket(&self, bucket: [BridgeLine; 3]) {
+        let mut ba_obj = self.ba.lock().unwrap();
+        ba_obj.add_spare_bucket(bucket);
+    }
+
     pub fn add_unreachable(&self, bridgeline: BridgeLine) -> bool {
         let mut ba_obj = self.ba.lock().unwrap();
         let mut db_obj = self.db.lock().unwrap();
diff --git a/crates/lox-distributor/src/main.rs b/crates/lox-distributor/src/main.rs
index 5b66f61..13d4e41 100644
--- a/crates/lox-distributor/src/main.rs
+++ b/crates/lox-distributor/src/main.rs
@@ -182,6 +182,7 @@ async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
     let context = lox_context::LoxServerContext {
         db: Arc::new(Mutex::new(bridgedb)),
         ba: Arc::new(Mutex::new(lox_auth)),
+        extra_bridges: Arc::new(Mutex::new(Vec::new())),
     };
 
     while let Some(cmd) = context_rx.recv().await {
@@ -190,16 +191,15 @@ async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
         match cmd {
             Rdsys { resourcediff } => {
                 if let Some(new_resources) = resourcediff.new {
+                    let mut count = 0;
+                    let mut bucket = [
+                        BridgeLine::default(),
+                        BridgeLine::default(),
+                        BridgeLine::default(),
+                    ];
                     for pt in new_resources {
                         println!("A NEW RESOURCE: {:?}", pt);
-                        let mut bucket = [
-                            BridgeLine::default(),
-                            BridgeLine::default(),
-                            BridgeLine::default(),
-                        ];
-                        let mut count = 0;
                         for resource in pt.1 {
-                        let test = resource.get_uid().unwrap();
                             let mut ip_bytes: [u8; 16] = [0; 16];
                             ip_bytes[..resource.address.len()]
                                 .copy_from_slice(resource.address.as_bytes());
@@ -229,6 +229,8 @@ async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
                                 bucket[count] = bridgeline;
                                 count += 1;
                             } else {
+                                // TODO: Decide the circumstances under which a bridge is allocated to an open_inv or spare bucket,
+                                // eventually also do some more fancy grouping of new resources, i.e., by type or region
                                 context.add_openinv_bucket(bucket);
                                 count = 0;
                                 bucket = [
@@ -239,6 +241,17 @@ async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
                             }
                         }
                     }
+                    // Handle the extra buckets that were not allocated already
+                    if count != 0 {
+                        for val in 0..count {
+                            if context.extra_bridges.lock().unwrap().len() < 2 {
+                                context.append_extra_bridges(bucket[val]);
+                            } else {
+                                bucket = context.remove_extra_bridges();
+                                context.add_spare_bucket(bucket);
+                            }
+                        }
+                    }
                 }
                 if let Some(changed_resources) = resourcediff.changed {
                     for pt in changed_resources {
@@ -274,7 +287,13 @@ async fn context_manager(mut context_rx: mpsc::Receiver<Command>) {
                             if res {
                                 println!("BridgeLine successfully updated: {:?}", bridgeline);
                             } else {
-                                println!("'Changed' BridgeLine NOT UPDATED!! : {:?}", bridgeline);
+                                println!("BridgeLine: {:?} not found in Lox's Bridgetable. Save it as a new resource for now!", bridgeline);
+                                if context.extra_bridges.lock().unwrap().len() < 2 {
+                                    context.append_extra_bridges(bridgeline);
+                                } else {
+                                    let bucket = context.remove_extra_bridges();
+                                    context.add_spare_bucket(bucket);
+                                }
                                 //TODO probably do something else here
                             }
                         }