Fix single step MFA locations without license

crates/defguard_core/src/grpc/client_version.rs

@@ -4,8 +4,6 @@ use defguard_proto::{client_types::ClientPlatformInfo, proxy::DeviceInfo};
 use prost::Message;
 use semver::Version;
 
-use crate::enterprise::is_business_license_active;
-
 const MOBILE_OS_FAMILIES: [&str; 2] = ["android", "ios"];
 
 fn is_mobile(platform: &ClientPlatformInfo) -> bool {
@@ -185,24 +183,21 @@ impl ClientFeature {
 }
 
 /// Returns `true` when a location should be omitted from a device's config because the location's
-/// MFA configuration has no legacy equivalent and either the device's client version does not
-/// support multi-step MFA or multi-step MFA is unavailable without an active business license.
+/// MFA configuration has no legacy equivalent and the device's client version does not support
+/// multi-step MFA.
+///
+/// License tiers are gated separately by [`crate::grpc::should_prevent_mfa_location_usage`].
 #[must_use]
 pub fn should_omit_location_for_device(
     location_mfa_mode: Option<LocationMfaMode>,
     device_info: Option<&DeviceInfo>,
 ) -> bool {
-    location_mfa_mode.is_none()
-        && (!ClientFeature::MultiStepMfa.is_supported_by_device(device_info)
-            || !is_business_license_active())
+    location_mfa_mode.is_none() && !ClientFeature::MultiStepMfa.is_supported_by_device(device_info)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::enterprise::license::{
-        License, LicenseTier, SupportType, get_cached_license, set_cached_license,
-    };
 
     // Helper function to create DeviceInfo
     fn create_device_info(
@@ -791,28 +786,8 @@ mod tests {
         );
     }
 
-    /// Builds a valid Business-tier license for tests that exercise licensed behavior.
-    ///
-    /// Setting this mutates the process-global license cache, so callers save and restore it
-    /// around their body. The restore is best-effort: a parallel test that also mutates the cache
-    /// can still race.
-    fn business_license() -> License {
-        License {
-            customer_id: "test".to_owned(),
-            subscription: false,
-            valid_until: None,
-            limits: None,
-            version_date_limit: None,
-            tier: LicenseTier::Business,
-            support_type: SupportType::Basic,
-            features: Vec::new(),
-        }
-    }
-
     #[test]
     fn test_should_omit_location_for_device() {
-        let saved_license = get_cached_license().clone();
-
         // Legacy client (below the MultiStepMfa floor).
         let legacy = create_device_info(
             Some("2.1.0".to_owned()),
@@ -844,14 +819,7 @@ mod tests {
         // Legacy client with no legacy equivalent is omitted.
         assert!(should_omit_location_for_device(None, Some(&legacy)));
 
-        // Capable client with no legacy equivalent but no business license is omitted.
-        set_cached_license(None);
-        assert!(should_omit_location_for_device(None, Some(&capable)));
-
-        // Capable client with no legacy equivalent and an active business license is included.
-        set_cached_license(Some(business_license()));
+        // Capable client with no legacy equivalent is included.
         assert!(!should_omit_location_for_device(None, Some(&capable)));
-
-        set_cached_license(saved_license);
     }
 }

crates/defguard_core/src/grpc/utils/tests.rs

@@ -334,6 +334,53 @@ async fn test_multi_step_location_omitted_for_legacy_client(
     set_cached_license(saved_license);
 }
 
+#[sqlx::test]
+#[allow(deprecated)]
+async fn test_single_step_location_sent_on_free_tier(_: PgPoolOptions, options: PgConnectOptions) {
+    let pool = setup_pool(options).await;
+    init_settings(&pool).await;
+    let saved_license = get_cached_license().clone();
+    set_cached_license(None);
+    let user = create_user(&pool).await;
+    let device = create_device(&pool, user.id).await;
+
+    // A TOTP-only step fits the Free tier but has no legacy equivalent.
+    let single_step = create_location(
+        &pool,
+        "single-step-location",
+        1,
+        vec![vec![VpnClientMfaMethod::Totp]],
+    )
+    .await;
+    let multi_step = create_location(
+        &pool,
+        "multi-step-location",
+        2,
+        vec![
+            vec![VpnClientMfaMethod::Totp],
+            vec![VpnClientMfaMethod::Email],
+        ],
+    )
+    .await;
+    attach_device(&pool, single_step.id, device.id).await;
+    attach_device(&pool, multi_step.id, device.id).await;
+
+    let response = build_device_config_response(&pool, device, None, Some(device_info("2.2.0")))
+        .await
+        .expect("failed to build config for capable client");
+    let names: Vec<&str> = response
+        .configs
+        .iter()
+        .map(|config| config.network_name.as_str())
+        .collect();
+    assert_eq!(names, vec!["single-step-location"]);
+    let config = &response.configs[0];
+    assert_eq!(config.location_mfa_mode, None);
+    assert_eq!(config.steps.len(), 1);
+
+    set_cached_license(saved_license);
+}
+
 #[sqlx::test]
 async fn test_network_device_rejected_on_mfa_locations(
     _: PgPoolOptions,

crates/defguard_proxy_manager/src/servers/enrollment.rs

@@ -34,6 +34,7 @@ use defguard_core::{
         client_version::{
             ClientFeature, client_registration_source, should_omit_location_for_device,
         },
+        should_prevent_mfa_location_usage,
         utils::{build_device_config_response, parse_client_ip_agent, to_wire_device_config},
     },
     handlers::user::check_password_strength,
@@ -949,8 +950,43 @@ impl EnrollmentServer {
             device.wireguard_pubkey, user.username, user.id,
         );
 
+        // Don't send locations whose MFA policy is above the active license.
+        let mut licensed_configs = Vec::with_capacity(configs.len());
+        for config in configs {
+            let location = WireguardNetwork::find_by_id(&mut *transaction, config.network_id)
+                .await
+                .map_err(|err| {
+                    error!(
+                        "Failed to fetch WireguardNetwork with ID {}: {err}",
+                        config.network_id
+                    );
+                    Status::internal("unexpected error")
+                })?
+                .ok_or_else(|| {
+                    error!("WireguardNetwork with ID {} not found", config.network_id);
+                    Status::internal("unexpected error")
+                })?;
+            if should_prevent_mfa_location_usage(&mut *transaction, &location)
+                .await
+                .map_err(|err| {
+                    error!(
+                        "Failed to read the MFA policy of location {}: {err}",
+                        location.name
+                    );
+                    Status::internal("unexpected error")
+                })?
+            {
+                warn!(
+                    "Tried to use location {} whose MFA policy is above the active licence.",
+                    location.name
+                );
+                continue;
+            }
+            licensed_configs.push(config);
+        }
+
         // Don't send them service locations or posture-checked locations if they don't support it
-        let configs = configs
+        let configs = licensed_configs
             .into_iter()
             .filter(|config| {
                 config.service_location_mode == ServiceLocationMode::Disabled

crates/defguard_proxy_manager/src/tests/proxy_manager/handler/enrollment.rs

@@ -26,15 +26,15 @@ use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
 use tokio::time::timeout;
 
 use super::support::{
-    STRONG_PASSWORD, assert_device_config_response, assert_error_response,
+    STRONG_PASSWORD, assert_device_config_response, assert_error_response, clear_test_license,
     complete_proxy_handshake, configure_ldap, configure_smtp, create_acl_network,
-    create_enrollment_token, create_mfa_network, create_multi_step_mfa_network, create_network,
-    create_polling_token, create_user, create_user_with_device, generate_totp_code,
-    insert_acl_rule_for_network, make_device_info, send_activate_user,
-    send_activate_user_without_password, send_code_mfa_setup_finish, send_code_mfa_setup_start,
-    send_mfa_config_authorize, send_mfa_config_start, set_test_license_business,
-    set_test_license_enterprise, setup_user_totp_mfa, start_enrollment_session,
-    totp_code_from_base32_secret,
+    create_enrollment_token, create_mfa_network, create_multi_step_mfa_network,
+    create_multi_step_mfa_network_with_steps, create_network, create_polling_token, create_user,
+    create_user_with_device, generate_totp_code, insert_acl_rule_for_network, make_device_info,
+    send_activate_user, send_activate_user_without_password, send_code_mfa_setup_finish,
+    send_code_mfa_setup_start, send_mfa_config_authorize, send_mfa_config_start,
+    set_test_license_business, set_test_license_enterprise, setup_user_totp_mfa,
+    start_enrollment_session, totp_code_from_base32_secret,
 };
 use crate::tests::common::{HandlerTestContext, TEST_TIMEOUT};
 
@@ -269,6 +269,54 @@ async fn test_new_device_omits_multi_step_location_for_legacy_client(
     context.finish().await.expect_server_finished().await;
 }
 
+/// Without a license, a `NewDevice` enrollment must send a single-step location to a capable
+/// client even when the step has no legacy equivalent, and must omit a multi-step location.
+#[sqlx::test]
+#[allow(deprecated)]
+async fn test_new_device_sends_single_step_location_on_free_tier(
+    _: PgPoolOptions,
+    options: PgConnectOptions,
+) {
+    let mut context = HandlerTestContext::new(options).await;
+    complete_proxy_handshake(&mut context).await;
+    clear_test_license();
+
+    let single_step =
+        create_multi_step_mfa_network_with_steps(&context.pool, vec![vec![MfaMethod::Totp.into()]])
+            .await;
+    create_multi_step_mfa_network(&context.pool).await;
+
+    let user = create_user(&context.pool).await;
+    let token = create_enrollment_token(&context.pool, user.id, Some(user.id)).await;
+    start_enrollment_session(&mut context, &token.id).await;
+    context.mock_proxy().send_request(CoreRequest {
+        id: 1,
+        device_info: Some(DeviceInfo {
+            version: Some("2.2.0".to_owned()),
+            ..make_device_info()
+        }),
+        payload: Some(core_request::Payload::NewDevice(NewDevice {
+            name: "Free Tier Laptop".to_owned(),
+            pubkey: "AA0aJzRBTltodYKPnKm2w9Dd6vcEER4rOEVSX2x5hpM=".to_owned(),
+            token: Some(token.id.clone()),
+        })),
+    });
+
+    let response = context.mock_proxy_mut().recv_outbound().await;
+    let cfg = assert_device_config_response(&response);
+    let names: Vec<&str> = cfg
+        .configs
+        .iter()
+        .map(|config| config.network_name.as_str())
+        .collect();
+    assert_eq!(names, vec![single_step.name.as_str()]);
+    let config = &cfg.configs[0];
+    assert_eq!(config.location_mfa_mode, None);
+    assert_eq!(config.steps.len(), 1);
+
+    context.finish().await.expect_server_finished().await;
+}
+
 /// Happy path: submit a strong password through an active enrollment session →
 /// handler returns `Empty`, the user's `enrollment_pending` flag is cleared, and
 /// the user gains a password hash in the DB.  A `BidiStreamEvent::Enrollment`