From 9a45de587412a95fefeb4e1dfc61c4505866e53f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Mon, 21 Oct 2024 13:15:08 +0300
Subject: [PATCH] Remove stale free riders after 24 hours (#5002)

---
 common/wireguard/src/error.rs       |  3 +++
 common/wireguard/src/peer_handle.rs | 26 +++++++++++++++++++-------
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/common/wireguard/src/error.rs b/common/wireguard/src/error.rs
index ebf4928f15..65a7269952 100644
--- a/common/wireguard/src/error.rs
+++ b/common/wireguard/src/error.rs
@@ -17,4 +17,7 @@ pub enum Error {
 
     #[error("{0}")]
     GatewayStorage(#[from] nym_gateway_storage::error::StorageError),
+
+    #[error("{0}")]
+    SystemTime(#[from] std::time::SystemTimeError),
 }
diff --git a/common/wireguard/src/peer_handle.rs b/common/wireguard/src/peer_handle.rs
index 81528320dc..cd91d99b3c 100644
--- a/common/wireguard/src/peer_handle.rs
+++ b/common/wireguard/src/peer_handle.rs
@@ -13,10 +13,12 @@ use nym_gateway_storage::Storage;
 use nym_task::TaskClient;
 use nym_wireguard_types::DEFAULT_PEER_TIMEOUT_CHECK;
 use std::sync::Arc;
+use std::time::{Duration, SystemTime};
 use tokio::sync::{mpsc, RwLock};
 use tokio_stream::{wrappers::IntervalStream, StreamExt};
 
 pub(crate) type SharedBandwidthStorageManager<St> = Arc<RwLock<BandwidthStorageManager<St>>>;
+const AUTO_REMOVE_AFTER: Duration = Duration::from_secs(60 * 60 * 24); // 24 hours
 
 pub struct PeerHandle<St> {
     storage: St,
@@ -26,6 +28,7 @@ pub struct PeerHandle<St> {
     request_tx: mpsc::Sender<PeerControlRequest>,
     timeout_check_interval: IntervalStream,
     task_client: TaskClient,
+    startup_timestamp: SystemTime,
 }
 
 impl<St: Storage + Clone + 'static> PeerHandle<St> {
@@ -50,14 +53,11 @@ impl<St: Storage + Clone + 'static> PeerHandle<St> {
             request_tx,
             timeout_check_interval,
             task_client,
+            startup_timestamp: SystemTime::now(),
         }
     }
 
-    async fn remove_depleted_peer(&self) -> Result<bool, Error> {
-        log::debug!(
-            "Peer {} doesn't have bandwidth anymore, removing it",
-            self.public_key.to_string()
-        );
+    async fn remove_peer(&self) -> Result<bool, Error> {
         let (response_tx, response_rx) = oneshot::channel();
         self.request_tx
             .send(PeerControlRequest::RemovePeer {
@@ -91,13 +91,25 @@ impl<St: Storage + Clone + 'static> PeerHandle<St> {
                 .await
                 .is_err()
             {
-                let success = self.remove_depleted_peer().await?;
+                let success = self.remove_peer().await?;
                 return Ok(!success);
             }
         } else {
+            if SystemTime::now().duration_since(self.startup_timestamp)? >= AUTO_REMOVE_AFTER {
+                log::debug!(
+                    "Peer {} has been present for 24 hours, removing it",
+                    self.public_key.to_string()
+                );
+                let success = self.remove_peer().await?;
+                return Ok(!success);
+            }
             let spent_bandwidth = kernel_peer.rx_bytes + kernel_peer.tx_bytes;
             if spent_bandwidth >= BANDWIDTH_CAP_PER_DAY {
-                let success = self.remove_depleted_peer().await?;
+                log::debug!(
+                    "Peer {} doesn't have bandwidth anymore, removing it",
+                    self.public_key.to_string()
+                );
+                let success = self.remove_peer().await?;
                 return Ok(!success);
             }
         }