From 89b7fdf921c88f8d22a474b2497aaa1d743d8663 Mon Sep 17 00:00:00 2001
From: Ian Goldberg <iang@uwaterloo.ca>
Date: Tue, 27 Apr 2021 18:26:49 -0400
Subject: [PATCH] Unit test for the bridge table

---
 crates/lox-library/Cargo.toml          |  1 +
 crates/lox-library/src/bridge_table.rs | 85 +++++++++++++++++++++++++-
 2 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/crates/lox-library/Cargo.toml b/crates/lox-library/Cargo.toml
index 388ae60..ee389cd 100644
--- a/crates/lox-library/Cargo.toml
+++ b/crates/lox-library/Cargo.toml
@@ -15,6 +15,7 @@ sha2 = "0.9"
 lazy_static = "1"
 hex_fmt = "0.3"
 aes-gcm = "0.8"
+base64 = "0.13"
 
 [features]
 default = ["u64_backend"]
diff --git a/crates/lox-library/src/bridge_table.rs b/crates/lox-library/src/bridge_table.rs
index 6116857..dfdc84f 100644
--- a/crates/lox-library/src/bridge_table.rs
+++ b/crates/lox-library/src/bridge_table.rs
@@ -14,7 +14,7 @@ use rand::RngCore;
 use std::convert::TryInto;
 
 /// Each bridge information line is serialized into this many bytes
-pub const BRIDGE_BYTES: usize = 128;
+pub const BRIDGE_BYTES: usize = 220;
 
 /// The max number of bridges per bucket
 pub const MAX_BRIDGES_PER_BUCKET: usize = 3;
@@ -32,8 +32,8 @@ pub struct BridgeLine {
     pub addr: [u8; 16],
     /// port
     pub port: u16,
-    /// other protocol information, including pluggable trasport, public
-    /// key, etc.
+    /// other protocol information, including pluggable transport,
+    /// public key, etc.
     pub info: [u8; BRIDGE_BYTES - 18],
 }
 
@@ -85,6 +85,40 @@ impl BridgeLine {
         }
         res
     }
+    /// Create a random BridgeLine for testing
+    #[cfg(test)]
+    pub fn random() -> Self {
+        let mut rng = rand::thread_rng();
+        let mut res: Self = Default::default();
+        // Pick a random 4-byte address
+        let mut addr: [u8; 4] = [0; 4];
+        rng.fill_bytes(&mut addr);
+        // If the leading byte is 224 or more, that's not a valid IPv4
+        // address.  Choose an IPv6 address instead (but don't worry too
+        // much about it being well formed).
+        if addr[0] >= 224 {
+            rng.fill_bytes(&mut res.addr);
+        } else {
+            // Store an IPv4 address as a v4-mapped IPv6 address
+            res.addr[10] = 255;
+            res.addr[11] = 255;
+            res.addr[12..16].copy_from_slice(&addr);
+        };
+        let ports: [u16; 4] = [443, 4433, 8080, 43079];
+        let portidx = (rng.next_u32() % 4) as usize;
+        res.port = ports[portidx];
+        let mut fingerprint: [u8; 20] = [0; 20];
+        let mut cert: [u8; 52] = [0; 52];
+        rng.fill_bytes(&mut fingerprint);
+        rng.fill_bytes(&mut cert);
+        let infostr: String = format!(
+            "obfs4 {} cert={} iat-mode=0",
+            hex_fmt::HexFmt(fingerprint),
+            base64::encode_config(cert, base64::STANDARD_NO_PAD)
+        );
+        res.info[..infostr.len()].copy_from_slice(infostr.as_bytes());
+        res
+    }
 }
 
 /// A BridgeTable is the internal structure holding the buckets
@@ -155,3 +189,48 @@ impl BridgeTable {
         ))
     }
 }
+
+// Unit tests that require access to private fields
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_bridge_table() -> Result<(), aead::Error> {
+        // Create an empty bridge table
+        let mut btable: BridgeTable = Default::default();
+        // Make 20 buckets with one random bridge each
+        for _ in 0..20 {
+            let bucket: [BridgeLine; 3] =
+                [BridgeLine::random(), Default::default(), Default::default()];
+            btable.new_bucket(bucket);
+        }
+        // And 20 more with three random bridges each
+        for _ in 0..20 {
+            let bucket: [BridgeLine; 3] = [
+                BridgeLine::random(),
+                BridgeLine::random(),
+                BridgeLine::random(),
+            ];
+            btable.new_bucket(bucket);
+        }
+        // Create the encrypted bridge table
+        btable.encrypt_table();
+        // Try to decrypt a 1-bridge bucket
+        let key7 = btable.keys[7];
+        let encbucket7 = btable.encbuckets[7];
+        let bucket7 = BridgeTable::decrypt_bucket(&key7, &encbucket7)?;
+        println!("bucket 7 = {:?}", bucket7);
+        // Try to decrypt a 3-bridge bucket
+        let key24 = btable.keys[24];
+        let encbucket24 = btable.encbuckets[24];
+        let bucket24 = BridgeTable::decrypt_bucket(&key24, &encbucket24)?;
+        println!("bucket 24 = {:?}", bucket24);
+        // Try to decrypt a bucket with the wrong key
+        let key12 = btable.keys[12];
+        let encbucket15 = btable.encbuckets[15];
+        let res = BridgeTable::decrypt_bucket(&key12, &encbucket15).unwrap_err();
+        println!("bucket key mismatch = {:?}", res);
+        Ok(())
+    }
+}