Rework the instance config fetching

src/enterprise/db/models/enterprise_settings.rs

@@ -4,7 +4,7 @@ use struct_patch::Patch;
 
 use crate::enterprise::license::{get_cached_license, validate_license};
 
-#[derive(Model, Deserialize, Serialize, Patch)]
+#[derive(Debug, Model, Deserialize, Serialize, Patch)]
 #[patch(attribute(derive(Serialize, Deserialize)))]
 pub struct EnterpriseSettings {
     #[serde(skip)]

src/enterprise/grpc/polling.rs

@@ -5,7 +5,7 @@ use crate::{
     enterprise::license::{get_cached_license, validate_license},
     grpc::{
         proto::{InstanceInfoRequest, InstanceInfoResponse},
-        utils::{build_device_config_response, build_instance_config_response},
+        utils::build_device_config_response,
     },
 };
 
@@ -26,7 +26,7 @@ impl PollingServer {
         // Polling service is enterprise-only, check the lincense
         if validate_license(get_cached_license().as_ref()).is_err() {
             debug!("No valid license, denying instance polling info");
-            return Err(Status::permission_denied("no valid license"));
+            return Err(Status::failed_precondition("no valid license"));
         }
 
         // Validate the token
@@ -83,10 +83,8 @@ impl PollingServer {
         // Build & return polling info
         let device_config =
             build_device_config_response(&self.pool, &device.wireguard_pubkey).await?;
-        let instance_config = build_instance_config_response(&self.pool).await?;
         Ok(InstanceInfoResponse {
             device_config: Some(device_config),
-            instance_config: Some(instance_config),
         })
     }
 }

src/grpc/enrollment.rs

@@ -126,10 +126,13 @@ impl EnrollmentServer {
                 );
                 return Err(Status::permission_denied("user is disabled"));
             };
-            info!("User {}({:?}) is active", user.username, user.id);
+            info!(
+                "User {}({:?}) is active, proceeding with enrollment",
+                user.username, user.id
+            );
 
-            let mut transaction = self.pool.begin().await.map_err(|_| {
-                error!("Failed to begin a transaction for enrollment.");
+            let mut transaction = self.pool.begin().await.map_err(|err| {
+                error!("Failed to begin a transaction for enrollment: {err}");
                 Status::internal("unexpected error")
             })?;
 
@@ -150,23 +153,36 @@ impl EnrollmentServer {
             );
 
             debug!(
-                "Retrieving settings for enrollment purpose for user {}({:?}).",
+                "Retrieving settings for enrollment of user {}({:?}).",
                 user.username, user.id
             );
             let settings = Settings::get_settings(&mut *transaction)
                 .await
-                .map_err(|_| {
-                    error!("Failed to get settings.");
+                .map_err(|err| {
+                    error!("Failed to get settings: {err}");
                     Status::internal("unexpected error")
                 })?;
             debug!("Settings: {settings:?}");
 
+            debug!(
+                "Retrieving enterprise settings for enrollment of user {}({:?}).",
+                user.username, user.id
+            );
+            let enterprise_settings =
+                EnterpriseSettings::get(&mut *transaction)
+                    .await
+                    .map_err(|err| {
+                        error!("Failed to get enterprise settings: {err}");
+                        Status::internal("unexpected error")
+                    })?;
+            debug!("Enterprise settings: {enterprise_settings:?}");
+
             let vpn_setup_optional = settings.enrollment_vpn_step_optional;
             debug!(
                 "Retrieving instance info for user {}({:?}).",
                 user.username, user.id
             );
-            let instance_info = InstanceInfo::new(settings, &user.username);
+            let instance_info = InstanceInfo::new(settings, &user.username, enterprise_settings);
             debug!("Instance info {instance_info:?}");
 
             debug!(
@@ -176,16 +192,16 @@ impl EnrollmentServer {
             let (username, user_id) = (user.username.clone(), user.id);
             let user_info = InitialUserInfo::from_user(&self.pool, user)
                 .await
-                .map_err(|_| {
+                .map_err(|err| {
                     error!(
-                        "Failed to get user info for user {}({:?})",
+                        "Failed to get user info for user {}({:?}): {err}",
                         username, user_id,
                     );
                     Status::internal("unexpected error")
                 })?;
             debug!("User info {user_info:?}");
 
-            debug!("Try to get basic admin info...");
+            debug!("Trying to get basic admin info...");
             let admin_info = admin.map(AdminInfo::from);
             debug!("Admin info {admin_info:?}");
 
@@ -196,8 +212,8 @@ impl EnrollmentServer {
             let enterprise_settings =
                 EnterpriseSettings::get(&mut *transaction)
                     .await
-                    .map_err(|_| {
-                        error!("Failed to get enterprise settings");
+                    .map_err(|err| {
+                        error!("Failed to get enterprise settings: {err}");
                         Status::internal("unexpected error")
                     })?;
             let enrollment_settings = super::proto::Settings {
@@ -216,8 +232,8 @@ impl EnrollmentServer {
             };
             debug!("Response {response:?}");
 
-            transaction.commit().await.map_err(|_| {
-                error!("Failed to commit transaction");
+            transaction.commit().await.map_err(|err| {
+                error!("Failed to commit transaction: {err}");
                 Status::internal("unexpected error")
             })?;
 
@@ -278,8 +294,8 @@ impl EnrollmentServer {
         }
         debug!("User is active.");
 
-        let mut transaction = self.pool.begin().await.map_err(|_| {
-            error!("Failed to begin transaction");
+        let mut transaction = self.pool.begin().await.map_err(|err| {
+            error!("Failed to begin transaction: {err}");
             Status::internal("unexpected error")
         })?;
 
@@ -303,8 +319,8 @@ impl EnrollmentServer {
         debug!("Retriving settings to send welcome email...");
         let settings = Settings::get_settings(&mut *transaction)
             .await
-            .map_err(|_| {
-                error!("Failed to get settings");
+            .map_err(|err| {
+                error!("Failed to get settings: {err}");
                 Status::internal("unexpected error")
             })?;
         debug!("Successfully retrived settings.");
@@ -339,8 +355,8 @@ impl EnrollmentServer {
             )?;
         }
 
-        transaction.commit().await.map_err(|_| {
-            error!("Failed to commit transaction");
+        transaction.commit().await.map_err(|err| {
+            error!("Failed to commit transaction: {err}");
             Status::internal("unexpected error")
         })?;
 
@@ -413,8 +429,11 @@ impl EnrollmentServer {
         );
         if let Some(device) = Device::find_by_pubkey(&self.pool, &request.pubkey)
             .await
-            .map_err(|_| {
-                error!("Failed to get device by its pubkey: {}", request.pubkey);
+            .map_err(|err| {
+                error!(
+                    "Failed to get device {} by its pubkey: {err}",
+                    request.pubkey
+                );
                 Status::internal("unexpected error")
             })?
         {
@@ -439,8 +458,8 @@ impl EnrollmentServer {
             user.username, user.id,
         );
 
-        let mut transaction = self.pool.begin().await.map_err(|_| {
-            error!("Failed to begin transaction");
+        let mut transaction = self.pool.begin().await.map_err(|err| {
+            error!("Failed to begin transaction: {err}");
             Status::internal("unexpected error")
         })?;
         device.save(&mut *transaction).await.map_err(|err| {
@@ -498,7 +517,22 @@ impl EnrollmentServer {
 );
                 Status::internal("unexpected error")
             })?;
-        debug!("Settings {settings:?}");
+        debug!("Settings: {settings:?}");
+
+        debug!(
+            "Fetching enterprise settings for device {} creation process for user {}({:?})",
+            device.wireguard_pubkey, user.username, user.id,
+        );
+        let enterprise_settings = EnterpriseSettings::get(&mut *transaction)
+            .await
+            .map_err(|err| {
+                error!(
+            "Failed to fetch enterprise settings for device {} creation process for user {}({:?}): {err}",
+            device.wireguard_pubkey, user.username, user.id,
+        );
+                Status::internal("unexpected error")
+            })?;
+        debug!("Enterprise settings: {enterprise_settings:?}");
 
         // create polling token for further client communication
         debug!(
@@ -560,7 +594,7 @@ impl EnrollmentServer {
         let response = DeviceConfigResponse {
             device: Some(device.into()),
             configs: configs.into_iter().map(Into::into).collect(),
-            instance: Some(InstanceInfo::new(settings, &user.username).into()),
+            instance: Some(InstanceInfo::new(settings, &user.username, enterprise_settings).into()),
             token: Some(token.token),
         };
         debug!("{response:?}.");

src/grpc/mod.rs

@@ -45,7 +45,11 @@ use self::{
 use crate::{
     auth::failed_login::FailedLoginMap,
     db::{AppEvent, Settings},
-    enterprise::grpc::polling::PollingServer,
+    enterprise::{
+        db::models::enterprise_settings::EnterpriseSettings,
+        grpc::polling::PollingServer,
+        license::{get_cached_license, validate_license},
+    },
     handlers::mail::send_gateway_disconnected_email,
     mail::Mail,
     server_config,
@@ -649,17 +653,25 @@ pub struct InstanceInfo {
     url: Url,
     proxy_url: Url,
     username: String,
+    disable_all_traffic: bool,
+    enterprise_enabled: bool,
 }
 
 impl InstanceInfo {
-    pub fn new<S: Into<String>>(settings: Settings, username: S) -> Self {
+    pub fn new<S: Into<String>>(
+        settings: Settings,
+        username: S,
+        enterprise_settings: EnterpriseSettings,
+    ) -> Self {
         let config = server_config();
         InstanceInfo {
             id: settings.uuid,
             name: settings.instance_name,
             url: config.url.clone(),
             proxy_url: config.enrollment_url.clone(),
             username: username.into(),
+            disable_all_traffic: enterprise_settings.disable_all_traffic,
+            enterprise_enabled: validate_license(get_cached_license().as_ref()).is_ok(),
         }
     }
 }
@@ -672,6 +684,8 @@ impl From<InstanceInfo> for crate::grpc::proto::InstanceInfo {
             url: instance.url.to_string(),
             proxy_url: instance.proxy_url.to_string(),
             username: instance.username,
+            disable_all_traffic: instance.disable_all_traffic,
+            enterprise_enabled: instance.enterprise_enabled,
         }
     }
 }

src/grpc/utils.rs

@@ -2,18 +2,15 @@ use ipnetwork::IpNetwork;
 use tonic::Status;
 
 use super::{
-    proto::{DeviceConfig as ProtoDeviceConfig, DeviceConfigResponse, InstanceConfigResponse},
+    proto::{DeviceConfig as ProtoDeviceConfig, DeviceConfigResponse},
     InstanceInfo,
 };
 use crate::{
     db::{
         models::{device::WireguardNetworkDevice, wireguard::WireguardNetwork},
         DbPool, Device, Settings, User,
     },
-    enterprise::{
-        db::models::enterprise_settings::EnterpriseSettings,
-        license::{get_cached_license, validate_license},
-    },
+    enterprise::db::models::enterprise_settings::EnterpriseSettings,
 };
 
 pub(crate) async fn build_device_config_response(
@@ -39,6 +36,11 @@ pub(crate) async fn build_device_config_response(
         Status::internal(format!("unexpected error: {err}"))
     })?;
 
+    let enterprise_settings = EnterpriseSettings::get(pool).await.map_err(|err| {
+        error!("Failed to get enterprise settings: {err}");
+        Status::internal(format!("unexpected error: {err}"))
+    })?;
+
     let mut configs: Vec<ProtoDeviceConfig> = Vec::new();
     let Some(device) = device else {
         return Err(Status::internal("device not found error"));
@@ -97,24 +99,7 @@ pub(crate) async fn build_device_config_response(
     Ok(DeviceConfigResponse {
         device: Some(device.into()),
         configs,
-        instance: Some(InstanceInfo::new(settings, &user.username).into()),
+        instance: Some(InstanceInfo::new(settings, &user.username, enterprise_settings).into()),
         token: None,
     })
 }
-
-pub(crate) async fn build_instance_config_response(
-    pool: &DbPool,
-) -> Result<InstanceConfigResponse, Status> {
-    debug!("Building instance config response");
-    let enterprise = validate_license(get_cached_license().as_ref()).is_ok();
-    let enterprise_settings = EnterpriseSettings::get(pool).await.map_err(|err| {
-        error!("Failed to get enterprise settings while building instance config response: {err}");
-        Status::internal("unexpected error")
-    })?;
-    debug!("Instance config response built");
-
-    Ok(InstanceConfigResponse {
-        enterprise,
-        disable_all_traffic: enterprise_settings.disable_all_traffic,
-    })
-}