fix VPN client disconnect event (#1234)

* refactor how clients are disconnected

* avoid marking inactive peers as connected

---------

Co-authored-by: Maciej Wójcik <[email protected]>

Author: wojcik91

crates/defguard_core/src/grpc/gateway/client_state.rs

@@ -166,33 +166,30 @@ impl ClientMap {
         let mut disconnected_clients = Vec::new();
 
         // get client state map for given location
-        let location_map = match self.0.get_mut(&location_id) {
-            Some(location_map) => location_map,
-            None => {
-                return Err(ClientMapError::LocationNotFound { location_id });
-            }
+        if let Some(location_map) = self.0.get_mut(&location_id) {
+            let disconnect_threshold = TimeDelta::seconds(peer_disconnect_threshold_secs.into());
+
+            // remove clients which have been inactive longer than given location's `peer_disconnect_threshold`
+            location_map.retain(|public_key, client_state| {
+                let now = Utc::now().naive_utc();
+                if (now - client_state.latest_handshake) > disconnect_threshold {
+                	debug!("VPN client's {public_key} ({}, ID {}) latest handshake ({}) was more than {peer_disconnect_threshold_secs} seconds ago. Marking VPN client as disconnected", client_state.device.name, client_state.device.id, client_state.latest_handshake);
+                    let disconnect_event_context = GrpcRequestContext::new(
+                        client_state.user_id,
+                        client_state.username.clone(),
+                        client_state.endpoint.ip(),
+                        client_state.device.id,
+                        client_state.device.name.clone(),
+                    );
+                    disconnected_clients
+                        .push((client_state.device.clone(), disconnect_event_context));
+
+                    return false;
+                };
+                true
+            });
         };
 
-        let disconnect_threshold = TimeDelta::seconds(peer_disconnect_threshold_secs.into());
-
-        // remove clients which have been inactive longer than given location's `peer_disconnect_threshold`
-        location_map.retain(|_public_key, client_state| {
-            let now = Utc::now().naive_utc();
-            if (now - client_state.latest_update) > disconnect_threshold {
-                let disconnect_event_context = GrpcRequestContext::new(
-                    client_state.user_id,
-                    client_state.username.clone(),
-                    client_state.endpoint.ip(),
-                    client_state.device.id,
-                    client_state.device.name.clone(),
-                );
-                disconnected_clients.push((client_state.device.clone(), disconnect_event_context));
-
-                return false;
-            };
-            true
-        });
-
         Ok(disconnected_clients)
     }
 }

crates/defguard_core/src/grpc/gateway/mod.rs

@@ -6,7 +6,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use chrono::{DateTime, Utc};
+use chrono::{DateTime, TimeDelta, Utc};
 use client_state::ClientMap;
 use sqlx::{query, Error as SqlxError, PgExecutor, PgPool};
 use thiserror::Error;
@@ -16,6 +16,7 @@ use tokio::{
         mpsc::{self, error::SendError, Receiver, UnboundedSender},
     },
     task::JoinHandle,
+    time::{interval, Duration},
 };
 use tokio_stream::Stream;
 use tonic::{metadata::MetadataMap, Code, Request, Response, Status};
@@ -34,6 +35,8 @@ use crate::{
     mail::Mail,
 };
 
+const PEER_DISCONNECT_INTERVAL: u64 = 60;
+
 /// Sends given `GatewayEvent` to be handled by gateway GRPC server
 ///
 /// If you want to use it inside the API context, use [`crate::AppState::send_wireguard_event`] instead
@@ -192,6 +195,7 @@ impl GatewayServer {
             .client_state
             .lock()
             .map_err(|_| GatewayServerError::ClientStateMutexError)?;
+        debug!("Current VPN client state map: {client_state:?}");
         Ok(client_state)
     }
 
@@ -726,8 +730,46 @@ impl gateway_service_server::GatewayService for GatewayServer {
         let network_id = Self::get_network_id(request.metadata())?;
         let gateway_hostname = Self::get_gateway_hostname(request.metadata())?;
         let mut stream = request.into_inner();
+        let mut disconnect_timer = interval(Duration::from_secs(PEER_DISCONNECT_INTERVAL));
+
+        loop {
+            // wait for a message or update client map at least once a mninute if no messages are received
+            let stats_update = tokio::select! {
+                message = stream.message() => {
+                    match message? {
+                        Some(update) => update,
+                        None => break, // Stream ended
+                    }
+                }
+                _ = disconnect_timer.tick() => {
+                    debug!("No stats updates received in last {PEER_DISCONNECT_INTERVAL} seconds. Updating disconnected VPN clients");
+                    // fetch location to get current peer disconnect threshold
+                    let location = self.fetch_location_from_db(network_id).await?;
+
+                    // perform client state operations in a dedicated block to drop mutex guard
+                    let disconnected_clients = {
+                        // acquire lock on client state map
+                        let mut client_map = self.get_client_state_guard()?;
+
+                        // disconnect inactive clients
+                        client_map.disconnect_inactive_vpn_clients_for_location(
+                            network_id,
+                            location.peer_disconnect_threshold,
+                        )?
+                    };
+
+                    // emit client disconnect events
+                    for (device, context) in disconnected_clients {
+                        self.emit_event(GrpcEvent::ClientDisconnected {
+                            context,
+                            location: location.clone(),
+                            device,
+                        })?;
+                    };
+                    continue;
+                }
+            };
 
-        while let Some(stats_update) = stream.message().await? {
             debug!("Received stats message: {stats_update:?}");
             let Some(stats_update::Payload::PeerStats(peer_stats)) = stats_update.payload else {
                 debug!("Received stats message is empty, skipping.");
@@ -781,30 +823,35 @@ impl gateway_service_server::GatewayService for GatewayServer {
                             );
                         }
                         None => {
-                            // mark new VPN client as connected
-                            client_map.connect_vpn_client(
-                                network_id,
-                                &gateway_hostname,
-                                &public_key,
-                                &device,
-                                &user,
-                                socket_addr,
-                                &stats,
-                            )?;
-
-                            // emit connection event
-                            let context = GrpcRequestContext::new(
-                                user.id,
-                                user.username.clone(),
-                                socket_addr.ip(),
-                                device.id,
-                                device.name.clone(),
-                            );
-                            self.emit_event(GrpcEvent::ClientConnected {
-                                context,
-                                location: location.clone(),
-                                device: device.clone(),
-                            })?;
+                            // don't mark inactive peers as connected
+                            if (Utc::now().naive_utc() - stats.latest_handshake)
+                                < TimeDelta::seconds(location.peer_disconnect_threshold.into())
+                            {
+                                // mark new VPN client as connected
+                                client_map.connect_vpn_client(
+                                    network_id,
+                                    &gateway_hostname,
+                                    &public_key,
+                                    &device,
+                                    &user,
+                                    socket_addr,
+                                    &stats,
+                                )?;
+
+                                // emit connection event
+                                let context = GrpcRequestContext::new(
+                                    user.id,
+                                    user.username.clone(),
+                                    socket_addr.ip(),
+                                    device.id,
+                                    device.name.clone(),
+                                );
+                                self.emit_event(GrpcEvent::ClientConnected {
+                                    context,
+                                    location: location.clone(),
+                                    device: device.clone(),
+                                })?;
+                            }
                         }
                     };