From 1edf8d8e353f90aa3c746e5442ac15341ff30d6e Mon Sep 17 00:00:00 2001
From: Rob <rob@espressosys.com>
Date: Mon, 18 Nov 2024 11:08:00 -0500
Subject: [PATCH] retroactively remove users from the whitelist

---
 cdn-broker/src/connections/mod.rs        |  6 ++++
 cdn-broker/src/lib.rs                    |  8 +++++
 cdn-broker/src/tasks/broker/mod.rs       |  1 +
 cdn-broker/src/tasks/broker/whitelist.rs | 45 ++++++++++++++++++++++++
 4 files changed, 60 insertions(+)
 create mode 100644 cdn-broker/src/tasks/broker/whitelist.rs

diff --git a/cdn-broker/src/connections/mod.rs b/cdn-broker/src/connections/mod.rs
index d4532c5..6499a99 100644
--- a/cdn-broker/src/connections/mod.rs
+++ b/cdn-broker/src/connections/mod.rs
@@ -242,6 +242,12 @@ impl Connections {
         self.brokers.keys().cloned().collect()
     }
 
+    /// Get all of the users that are connected to us. We use this when we need
+    /// to check if they are still whitelisted.
+    pub fn all_users(&self) -> Vec<UserPublicKey> {
+        self.users.keys().cloned().collect()
+    }
+
     /// Insert a broker with its connection into our map.
     pub fn add_broker(
         &mut self,
diff --git a/cdn-broker/src/lib.rs b/cdn-broker/src/lib.rs
index 87f375c..141aaf0 100644
--- a/cdn-broker/src/lib.rs
+++ b/cdn-broker/src/lib.rs
@@ -276,6 +276,11 @@ impl<R: RunDef> Broker<R> {
         let inner_ = self.inner.clone();
         let sync_task = AbortOnDropHandle(spawn(inner_.run_sync_task()));
 
+        // Spawn the whitelist task, which retroactively checks if existing users are still
+        // whitelisted
+        let inner_ = self.inner.clone();
+        let whitelist_task = AbortOnDropHandle(spawn(inner_.run_whitelist_task()));
+
         // Spawn the public (user) listener task
         // TODO: maybe macro this, since it's repeat code with the private listener task
         let inner_ = self.inner.clone();
@@ -307,6 +312,9 @@ impl<R: RunDef> Broker<R> {
             _ = broker_listener_task => {
                 Err(Error::Exited("broker listener task exited!".to_string()))
             }
+            _ = whitelist_task => {
+                Err(Error::Exited("whitelist task exited!".to_string()))
+            }
         }
     }
 }
diff --git a/cdn-broker/src/tasks/broker/mod.rs b/cdn-broker/src/tasks/broker/mod.rs
index 8c79717..c0954a0 100644
--- a/cdn-broker/src/tasks/broker/mod.rs
+++ b/cdn-broker/src/tasks/broker/mod.rs
@@ -9,3 +9,4 @@ pub mod heartbeat;
 pub mod listener;
 pub mod sender;
 pub mod sync;
+pub mod whitelist;
diff --git a/cdn-broker/src/tasks/broker/whitelist.rs b/cdn-broker/src/tasks/broker/whitelist.rs
new file mode 100644
index 0000000..4fe1437
--- /dev/null
+++ b/cdn-broker/src/tasks/broker/whitelist.rs
@@ -0,0 +1,45 @@
+// Copyright (c) 2024 Espresso Systems (espressosys.com)
+// This file is part of the Push-CDN repository.
+
+// You should have received a copy of the MIT License
+// along with the Push-CDN repository. If not, see <https://mit-license.org/>.
+
+//! The sync task syncs both users and topics to other brokers.
+
+use std::{sync::Arc, time::Duration};
+
+use cdn_proto::{def::RunDef, discovery::DiscoveryClient};
+use tokio::time::sleep;
+
+use crate::Inner;
+
+impl<Def: RunDef> Inner<Def> {
+    /// Run the whitelist task. This is responsible for checking if users are still whitelisted
+    /// and kicking them off the network if they are not.
+    pub async fn run_whitelist_task(self: Arc<Self>) {
+        // Clone the discovery client because it's behind an `Arc`
+        let mut discovery_client = self.discovery_client.clone();
+
+        loop {
+            // Run every minute
+            sleep(Duration::from_secs(60)).await;
+
+            // Get a list of all users
+            let users = self.connections.read().all_users();
+
+            // Make sure each user is still whitelisted
+            for user in users {
+                if !discovery_client
+                    .check_whitelist(&user)
+                    .await
+                    .unwrap_or(true)
+                {
+                    // Kick the user off the network if they are not
+                    self.connections
+                        .write()
+                        .remove_user(user, "not in whitelist");
+                }
+            }
+        }
+    }
+}