feat: desktop client MFA

src/db/models/device.rs

@@ -491,7 +491,6 @@ impl Device {
     pub async fn add_to_all_networks(
         &self,
         transaction: &mut PgConnection,
-        admin_group_name: &str,
     ) -> Result<(Vec<DeviceNetworkInfo>, Vec<DeviceConfig>), DeviceError> {
         info!("Adding device {} to all existing networks", self.name);
         let networks = WireguardNetwork::all(&mut *transaction).await?;
@@ -528,7 +527,7 @@ impl Device {
             }
 
             if let Ok(wireguard_network_device) = network
-                .add_device_to_network(&mut *transaction, self, admin_group_name, None)
+                .add_device_to_network(&mut *transaction, self, None)
                 .await
             {
                 debug!(

src/db/models/group.rs

@@ -2,6 +2,7 @@ use model_derive::Model;
 use sqlx::{query, query_as, query_scalar, Error as SqlxError, PgConnection, PgExecutor};
 
 use crate::db::{models::error::ModelError, User, WireguardNetwork};
+use crate::SERVER_CONFIG;
 
 #[derive(Model)]
 pub struct Group {
@@ -97,9 +98,12 @@ impl WireguardNetwork {
     pub async fn get_allowed_groups(
         &self,
         transaction: &mut PgConnection,
-        admin_group_name: &str,
     ) -> Result<Option<Vec<String>>, ModelError> {
         debug!("Returning a list of allowed groups for network {self}");
+        let admin_group_name = &SERVER_CONFIG
+            .get()
+            .expect("defguard config not found")
+            .admin_groupname;
         // get allowed groups from DB
         let mut groups = self.fetch_allowed_groups(&mut *transaction).await?;
 

src/db/models/wireguard.rs

@@ -302,11 +302,10 @@ impl WireguardNetwork {
     async fn get_allowed_devices(
         &self,
         transaction: &mut PgConnection,
-        admin_group_name: &str,
     ) -> Result<Vec<Device>, ModelError> {
         debug!("Fetching all allowed devices for network {}", self);
         let devices = match self
-            .get_allowed_groups(&mut *transaction, admin_group_name)
+            .get_allowed_groups(&mut *transaction)
             .await? {
             // devices need to be filtered by allowed group
             Some(allowed_groups) => {
@@ -338,15 +337,12 @@ impl WireguardNetwork {
     pub async fn add_all_allowed_devices(
         &self,
         transaction: &mut PgConnection,
-        admin_group_name: &str,
     ) -> Result<(), ModelError> {
         info!(
             "Assigning IPs in network {} for all existing devices ",
             self
         );
-        let devices = self
-            .get_allowed_devices(&mut *transaction, admin_group_name)
-            .await?;
+        let devices = self.get_allowed_devices(&mut *transaction).await?;
         for device in devices {
             device
                 .assign_network_ip(&mut *transaction, self, None)
@@ -360,13 +356,10 @@ impl WireguardNetwork {
         &self,
         transaction: &mut PgConnection,
         device: &Device,
-        admin_group_name: &str,
         reserved_ips: Option<&[IpAddr]>,
     ) -> Result<WireguardNetworkDevice, WireguardNetworkError> {
         info!("Assigning IP in network {self} for {device}");
-        let allowed_devices = self
-            .get_allowed_devices(&mut *transaction, admin_group_name)
-            .await?;
+        let allowed_devices = self.get_allowed_devices(&mut *transaction).await?;
         let allowed_device_ids: Vec<i64> =
             allowed_devices.iter().filter_map(|dev| dev.id).collect();
         if allowed_device_ids.contains(&device.get_id()?) {
@@ -389,14 +382,11 @@ impl WireguardNetwork {
     pub async fn sync_allowed_devices(
         &self,
         transaction: &mut PgConnection,
-        admin_group_name: &str,
         reserved_ips: Option<&[IpAddr]>,
     ) -> Result<Vec<GatewayEvent>, WireguardNetworkError> {
         info!("Synchronizing IPs in network {self} for all allowed devices ");
         // list all allowed devices
-        let allowed_devices = self
-            .get_allowed_devices(&mut *transaction, admin_group_name)
-            .await?;
+        let allowed_devices = self.get_allowed_devices(&mut *transaction).await?;
         // convert to a map for easier processing
         let mut allowed_devices: HashMap<i64, Device> = allowed_devices
             .into_iter()
@@ -485,12 +475,9 @@ impl WireguardNetwork {
         &self,
         transaction: &mut PgConnection,
         imported_devices: Vec<ImportedDevice>,
-        admin_group_name: &str,
     ) -> Result<(Vec<ImportedDevice>, Vec<GatewayEvent>), WireguardNetworkError> {
         let network_id = self.get_id()?;
-        let allowed_devices = self
-            .get_allowed_devices(&mut *transaction, admin_group_name)
-            .await?;
+        let allowed_devices = self.get_allowed_devices(&mut *transaction).await?;
         // convert to a map for easier processing
         let allowed_devices: HashMap<i64, Device> = allowed_devices
             .into_iter()
@@ -551,14 +538,11 @@ impl WireguardNetwork {
         &self,
         transaction: &mut PgConnection,
         mapped_devices: Vec<MappedDevice>,
-        admin_group_name: &str,
     ) -> Result<Vec<GatewayEvent>, WireguardNetworkError> {
         info!("Mapping user devices for network {}", self);
         let network_id = self.get_id()?;
         // get allowed groups for network
-        let allowed_groups = self
-            .get_allowed_groups(&mut *transaction, admin_group_name)
-            .await?;
+        let allowed_groups = self.get_allowed_groups(&mut *transaction).await?;
 
         let mut events = Vec::new();
         // use a helper hashmap to avoid repeated queries
@@ -629,9 +613,8 @@ impl WireguardNetwork {
             }
 
             // assign IPs in other networks
-            let (mut all_network_info, _configs) = device
-                .add_to_all_networks(&mut *transaction, admin_group_name)
-                .await?;
+            let (mut all_network_info, _configs) =
+                device.add_to_all_networks(&mut *transaction).await?;
 
             network_info.append(&mut all_network_info);
 

src/grpc/desktop_client_mfa.rs

@@ -0,0 +1,230 @@
+use crate::{
+    db::{
+        models::device::{DeviceInfo, DeviceNetworkInfo, WireguardNetworkDevice},
+        DbPool, Device, GatewayEvent, User, UserInfo, WireguardNetwork,
+    },
+    handlers::mail::send_email_mfa_code_email,
+    mail::Mail,
+};
+use std::collections::HashMap;
+use tokio::sync::{broadcast::Sender, mpsc::UnboundedSender};
+use tonic::Status;
+use uuid::Uuid;
+
+use super::proto::{
+    ClientMfaFinishRequest, ClientMfaFinishResponse, ClientMfaStartRequest, ClientMfaStartResponse,
+    MfaMethod,
+};
+
+struct ClientLoginSession {
+    pub location: WireguardNetwork,
+    pub device: Device,
+    pub user: User,
+}
+
+pub(super) struct ClientMfaServer {
+    pool: DbPool,
+    mail_tx: UnboundedSender<Mail>,
+    wireguard_tx: Sender<GatewayEvent>,
+    sessions: HashMap<String, ClientLoginSession>,
+}
+
+impl ClientMfaServer {
+    #[must_use]
+    pub fn new(
+        pool: DbPool,
+        mail_tx: UnboundedSender<Mail>,
+        wireguard_tx: Sender<GatewayEvent>,
+    ) -> Self {
+        Self {
+            pool,
+            mail_tx,
+            wireguard_tx,
+            sessions: HashMap::new(),
+        }
+    }
+
+    pub async fn start_client_mfa_login(
+        &mut self,
+        request: ClientMfaStartRequest,
+    ) -> Result<ClientMfaStartResponse, Status> {
+        info!("Starting desktop client login: {request:?}");
+        // fetch location
+        let Ok(Some(location)) =
+            WireguardNetwork::find_by_id(&self.pool, request.location_id).await
+        else {
+            error!("Failed to find location with ID {}", request.location_id);
+            return Err(Status::invalid_argument("location not found"));
+        };
+
+        // fetch device
+        let Ok(Some(device)) = Device::find_by_pubkey(&self.pool, &request.pubkey).await else {
+            error!("Failed to find device with pubkey {}", request.pubkey);
+            return Err(Status::invalid_argument("device not found"));
+        };
+
+        // fetch user
+        let Ok(Some(user)) = User::find_by_id(&self.pool, device.user_id).await else {
+            error!("Failed to find user with ID {}", device.user_id);
+            return Err(Status::invalid_argument("user not found"));
+        };
+        let user_info = UserInfo::from_user(&self.pool, &user).await.map_err(|_| {
+            error!("Failed to fetch user info for {}", user.username);
+            Status::internal("unexpected error")
+        })?;
+
+        // validate user is allowed to connect to a given location
+        let mut transaction = self.pool.begin().await.map_err(|_| {
+            error!("Failed to begin transaction");
+            Status::internal("unexpected error")
+        })?;
+        let allowed_groups = location
+            .get_allowed_groups(&mut transaction)
+            .await
+            .map_err(|err| {
+                error!("Failed to fetch allowed groups for location {location}: {err:?}");
+                Status::internal("unexpected error")
+            })?;
+        if let Some(groups) = allowed_groups {
+            // check if user belongs to one of allowed groups
+            if !groups
+                .iter()
+                .any(|allowed_group| user_info.groups.contains(allowed_group))
+            {
+                error!(
+                    "User {} not allowed to connect to location {location}",
+                    user.username
+                );
+                return Err(Status::unauthenticated("unauthorized"));
+            }
+        }
+
+        // check if selected method is enabled
+        match MfaMethod::try_from(request.method) {
+            Ok(MfaMethod::Totp) => {
+                if !user.totp_enabled {
+                    error!("TOTP not enabled for user {}", user.username);
+                    return Err(Status::invalid_argument(
+                        "selected MFA method not available",
+                    ));
+                }
+            }
+            Ok(MfaMethod::Email) => {
+                if !user.email_mfa_enabled {
+                    error!("Email MFA not enabled for user {}", user.username);
+                    return Err(Status::invalid_argument(
+                        "selected MFA method not available",
+                    ));
+                }
+                // send email code
+                send_email_mfa_code_email(&user, &self.mail_tx, None).map_err(|err| {
+                    error!(
+                        "Failed to send email MFA code for user {}: {err:?}",
+                        user.username
+                    );
+                    Status::internal("unexpected error")
+                })?;
+            }
+            Err(err) => {
+                error!("Invalid MFA method selected: {err}");
+                return Err(Status::invalid_argument("invalid MFA method selected"));
+            }
+        }
+
+        // generate auth token
+        let token = Uuid::new_v4().to_string();
+
+        // store login session
+        self.sessions.insert(
+            token.clone(),
+            ClientLoginSession {
+                location,
+                device,
+                user,
+            },
+        );
+
+        Ok(ClientMfaStartResponse { token })
+    }
+
+    pub async fn finish_client_mfa_login(
+        &mut self,
+        request: ClientMfaFinishRequest,
+    ) -> Result<ClientMfaFinishResponse, Status> {
+        info!("Finishing desktop client login: {request:?}");
+        // fetch login session
+        let Some(session) = self.sessions.remove(&request.token) else {
+            error!("Client login session not found");
+            return Err(Status::invalid_argument("login session not found"));
+        };
+        let device = session.device;
+        let location = session.location;
+        let user = session.user;
+
+        // validate email code
+        if !user.verify_email_mfa_code(request.code) {
+            error!("Provided email code is not valid");
+            return Err(Status::unauthenticated("unauthorized"));
+        };
+
+        // begin transaction
+        let mut transaction = self.pool.begin().await.map_err(|_| {
+            error!("Failed to begin transaction");
+            Status::internal("unexpected error")
+        })?;
+
+        // fetch device config for the location
+        let Ok(Some(mut network_device)) = WireguardNetworkDevice::find(
+            &mut *transaction,
+            device.id.expect("Missing device ID"),
+            location.id.expect("Missing location ID"),
+        )
+        .await
+        else {
+            error!("Failed to fetch network config for device {device} and location {location}");
+            return Err(Status::internal("unexpected error"));
+        };
+
+        // generate PSK
+        let key = WireguardNetwork::genkey();
+        network_device.preshared_key = Some(key.public.clone());
+
+        // authorize device for given location
+        network_device.is_authorized = true;
+
+        // save updated network config
+        network_device
+            .update(&mut *transaction)
+            .await
+            .map_err(|err| {
+                error!("Failed to update device network config {network_device:?}: {err:?}");
+                Status::internal("unexpected error")
+            })?;
+
+        // send gateway event
+        debug!("Sending `peer_create` message to gateway");
+        let device_info = DeviceInfo {
+            device,
+            network_info: vec![DeviceNetworkInfo {
+                network_id: location.id.expect("Missing location ID"),
+                device_wireguard_ip: network_device.wireguard_ip,
+                preshared_key: network_device.preshared_key,
+            }],
+        };
+        let event = GatewayEvent::DeviceCreated(device_info);
+        self.wireguard_tx.send(event).map_err(|err| {
+            error!("Error sending WireGuard event: {err}");
+            Status::internal("unexpected error")
+        })?;
+
+        // commit transaction
+        transaction.commit().await.map_err(|_| {
+            error!("Failed to commit transaction");
+            Status::internal("unexpected error")
+        })?;
+
+        Ok(ClientMfaFinishResponse {
+            preshared_key: key.public,
+        })
+    }
+}