From a32cbd44b0e009c6cd9677a1738f6b48fb3bd6bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Tue, 11 Jun 2024 10:05:38 +0000
Subject: [PATCH] Remove stale peers

---
 Cargo.lock                              |  1 +
 common/wireguard/Cargo.toml             |  1 +
 common/wireguard/src/peer_controller.rs | 48 +++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b06aeeec7e..702e67160b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5783,6 +5783,7 @@ dependencies = [
  "nym-task",
  "nym-wireguard-types",
  "tokio",
+ "tokio-stream",
  "x25519-dalek",
 ]
 
diff --git a/common/wireguard/Cargo.toml b/common/wireguard/Cargo.toml
index 50ea7573b6..9097c90ae8 100644
--- a/common/wireguard/Cargo.toml
+++ b/common/wireguard/Cargo.toml
@@ -25,3 +25,4 @@ nym-network-defaults = { path = "../network-defaults" }
 nym-task = { path = "../task" }
 nym-wireguard-types = { path = "../wireguard-types" }
 tokio = { workspace = true, features = ["rt-multi-thread", "net", "io-util"] }
+tokio-stream = { workspace = true }
diff --git a/common/wireguard/src/peer_controller.rs b/common/wireguard/src/peer_controller.rs
index a1a3334c3a..a325cf29db 100644
--- a/common/wireguard/src/peer_controller.rs
+++ b/common/wireguard/src/peer_controller.rs
@@ -1,13 +1,24 @@
 // Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use std::sync::Arc;
+use std::{
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
 
-use defguard_wireguard_rs::{host::Peer, key::Key, WireguardInterfaceApi};
+use defguard_wireguard_rs::{
+    host::{Host, Peer},
+    key::Key,
+    WGApi, WireguardInterfaceApi,
+};
 use tokio::sync::mpsc;
+use tokio_stream::{wrappers::IntervalStream, StreamExt};
 
 use crate::WgApiWrapper;
 
+const DEFAULT_PEER_TIMEOUT: Duration = Duration::from_secs(60 * 60); // 1 hour
+const DEFAULT_PEER_TIMEOUT_CHECK: Duration = Duration::from_secs(60); // 1 minute
+
 pub enum PeerControlMessage {
     AddPeer(Peer),
     RemovePeer(Key),
@@ -16,6 +27,7 @@ pub enum PeerControlMessage {
 pub struct PeerController {
     peer_rx: mpsc::UnboundedReceiver<PeerControlMessage>,
     wg_api: Arc<WgApiWrapper>,
+    timeout_check_interval: IntervalStream,
 }
 
 impl PeerController {
@@ -23,12 +35,42 @@ impl PeerController {
         wg_api: Arc<WgApiWrapper>,
         peer_rx: mpsc::UnboundedReceiver<PeerControlMessage>,
     ) -> Self {
-        PeerController { wg_api, peer_rx }
+        let timeout_check_interval = tokio_stream::wrappers::IntervalStream::new(
+            tokio::time::interval(DEFAULT_PEER_TIMEOUT_CHECK),
+        );
+        PeerController {
+            wg_api,
+            peer_rx,
+            timeout_check_interval,
+        }
+    }
+
+    fn remove_stale_peers(wg_api: &WGApi, host: Host) {
+        let current_timestamp = SystemTime::now();
+        for (key, peer) in host.peers.iter() {
+            if let Some(timestamp) = peer.last_handshake {
+                if let Ok(duration_since_handshake) = current_timestamp.duration_since(timestamp) {
+                    if duration_since_handshake > DEFAULT_PEER_TIMEOUT {
+                        if let Err(e) = wg_api.remove_peer(key) {
+                            log::error!("Could not remove stale peer: {:?}", e);
+                        } else {
+                            log::debug!("Removed stale peer {:?}", key);
+                        }
+                    }
+                }
+            }
+        }
     }
 
     pub async fn run(&mut self, mut task_client: nym_task::TaskClient) {
         loop {
             tokio::select! {
+                _ = self.timeout_check_interval.next() => {
+                    match self.wg_api.inner.read_interface_data() {
+                        Ok(host) => Self::remove_stale_peers(&self.wg_api.inner, host),
+                        Err(e) => { log::error!("Could not read peer data: {:?}", e); },
+                    }
+                }
                 _ = task_client.recv() => {
                     log::trace!("PeerController handler: Received shutdown");
                     break;