Posture assignment updates gateways

crates/defguard_core/src/enterprise/handlers/device_posture.rs

@@ -8,7 +8,7 @@ use axum::{
 use axum_extra::extract::Query as AxumExtraQuery;
 use defguard_common::db::{Id, NoId, models::WireguardNetwork};
 use serde::{Deserialize, Serialize};
-use sqlx::{Postgres, QueryBuilder};
+use sqlx::{PgConnection, Postgres, QueryBuilder};
 use utoipa::ToSchema;
 
 use crate::{
@@ -19,6 +19,7 @@ use crate::{
             DevicePosture, DevicePostureLocation, DevicePostureOsRule, DevicePostureSnapshot,
             OsType,
         },
+        firewall::try_get_location_firewall_config,
         handlers::{DevicePostureFeature, LicenseGated},
         posture::version_list::{
             ANDROID_OS_VERSIONS, CLIENT_VERSIONS, IOS_OS_VERSIONS, LINUX_KERNEL_VERSIONS,
@@ -27,16 +28,52 @@ use crate::{
     },
     error::WebError,
     events::{ApiEvent, ApiEventType, ApiRequestContext},
+    grpc::GatewayCommand,
     handlers::{
         ApiResponse, ApiResult,
         pagination::{PaginatedApiResponse, PaginatedApiResult, PaginationParams},
     },
+    location_management::allowed_peers::get_location_allowed_peers,
 };
 
 fn owned_client_versions(values: &[&str]) -> Vec<String> {
     values.iter().map(|value| (*value).to_owned()).collect()
 }
 
+fn same_id_set(left: &[Id], right: &[Id]) -> bool {
+    left.iter().copied().collect::<HashSet<_>>() == right.iter().copied().collect::<HashSet<_>>()
+}
+
+async fn build_location_peer_refresh_commands(
+    conn: &mut PgConnection,
+    location_ids: impl IntoIterator<Item = Id>,
+) -> Result<Vec<GatewayCommand>, WebError> {
+    let mut seen = HashSet::new();
+    let mut commands = Vec::new();
+
+    for location_id in location_ids {
+        if !seen.insert(location_id) {
+            continue;
+        }
+
+        let Some(location) = WireguardNetwork::find_by_id(&mut *conn, location_id).await? else {
+            warn!("Skipping gateway peer refresh for missing location {location_id}");
+            continue;
+        };
+
+        let peers = get_location_allowed_peers(&location, &mut *conn).await?;
+        let maybe_firewall_config = try_get_location_firewall_config(&location, &mut *conn).await?;
+        commands.push(GatewayCommand::NetworkModified(
+            location.id,
+            location,
+            peers,
+            maybe_firewall_config,
+        ));
+    }
+
+    Ok(commands)
+}
+
 /// Per-OS rule included in a posture check policy API.
 ///
 /// Adding this layer on top of the shared DB type allows us
@@ -939,7 +976,13 @@ pub async fn delete_device_posture(
     let os_rules = DevicePostureOsRule::find_by_posture(&appstate.pool, id).await?;
     let location_ids = DevicePostureLocation::find_by_posture(&appstate.pool, id).await?;
 
-    device_posture.clone().delete(&appstate.pool).await?;
+    let mut tx = appstate.pool.begin().await?;
+    device_posture.clone().delete(&mut *tx).await?;
+    let gateway_commands =
+        build_location_peer_refresh_commands(&mut tx, location_ids.clone()).await?;
+    tx.commit().await?;
+
+    appstate.send_multiple_gateway_commands(gateway_commands);
 
     appstate.emit_event(ApiEvent {
         context,
@@ -1101,10 +1144,18 @@ pub async fn set_postures_for_location(
     }
 
     let mut tx = appstate.pool.begin().await?;
+    let old_postures = DevicePostureLocation::find_by_location(&mut *tx, location_id).await?;
     let result =
         DevicePostureLocation::set_for_location(&mut tx, location_id, &data.postures).await?;
+    let gateway_commands = if same_id_set(&old_postures, &result) {
+        Vec::new()
+    } else {
+        build_location_peer_refresh_commands(&mut tx, [location_id]).await?
+    };
     tx.commit().await?;
 
+    appstate.send_multiple_gateway_commands(gateway_commands);
+
     appstate.emit_event(ApiEvent {
         context,
         event: Box::new(ApiEventType::LocationPosturesAssigned {
@@ -1168,10 +1219,22 @@ pub async fn set_locations_for_posture(
     }
 
     let mut tx = appstate.pool.begin().await?;
+    let old_locations = DevicePostureLocation::find_by_posture(&mut *tx, posture_id).await?;
     let result =
         DevicePostureLocation::set_for_posture(&mut tx, posture_id, &data.locations).await?;
+    let gateway_commands = if same_id_set(&old_locations, &result) {
+        Vec::new()
+    } else {
+        build_location_peer_refresh_commands(
+            &mut tx,
+            old_locations.iter().copied().chain(result.iter().copied()),
+        )
+        .await?
+    };
     tx.commit().await?;
 
+    appstate.send_multiple_gateway_commands(gateway_commands);
+
     appstate.emit_event(ApiEvent {
         context,
         event: Box::new(ApiEventType::DevicePostureLocationsAssigned {

crates/defguard_core/tests/integration/api/device_posture.rs

@@ -1,4 +1,4 @@
-use defguard_common::db::setup_pool;
+use defguard_common::db::{Id, setup_pool};
 use defguard_core::{
     enterprise::{
         db::models::device_posture::{DevicePosture, DevicePostureSnapshot},
@@ -13,6 +13,7 @@ use defguard_core::{
         },
     },
     events::ApiEventType,
+    grpc::GatewayCommand,
 };
 use reqwest::StatusCode;
 use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
@@ -45,6 +46,51 @@ async fn setup(options: PgConnectOptions) -> (TestClient, ClientState) {
     (client, state)
 }
 
+fn drain_gateway_events(gateway_rx: &mut tokio::sync::broadcast::Receiver<GatewayCommand>) {
+    while gateway_rx.try_recv().is_ok() {}
+}
+
+fn expect_network_modified_peers(
+    gateway_rx: &mut tokio::sync::broadcast::Receiver<GatewayCommand>,
+    location_id: Id,
+    expected_pubkeys: &[&str],
+) {
+    let event = gateway_rx.try_recv().unwrap();
+    let GatewayCommand::NetworkModified(id, _, peers, _) = event else {
+        panic!("expected NetworkModified, got {event:?}");
+    };
+    assert_eq!(id, location_id);
+    let actual_pubkeys = peers
+        .iter()
+        .map(|peer| peer.pubkey.as_str())
+        .collect::<Vec<_>>();
+    assert_eq!(actual_pubkeys, expected_pubkeys);
+}
+
+async fn create_posture(client: &TestClient, name: &str) -> ApiDevicePosture {
+    let response = client
+        .post("/api/v1/device-posture")
+        .json(&make_edit(name))
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::CREATED);
+    response.json().await
+}
+
+async fn create_admin_device(client: &TestClient) -> &'static str {
+    const PUBKEY: &str = "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=";
+    let response = client
+        .post("/api/v1/device/admin")
+        .json(&serde_json::json!({
+            "name": "posture-test-device",
+            "wireguard_pubkey": PUBKEY,
+        }))
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::CREATED);
+    PUBKEY
+}
+
 #[sqlx::test]
 async fn test_device_posture_enterprise_license_required(
     _: PgPoolOptions,
@@ -947,6 +993,167 @@ async fn test_device_posture_set_postures_for_location(
     }
 }
 
+#[sqlx::test]
+async fn test_assigning_first_posture_refreshes_gateway_with_no_direct_peers(
+    _: PgPoolOptions,
+    options: PgConnectOptions,
+) {
+    let (mut client, state) = setup(options).await;
+    let mut gateway_rx = state.gateway_rx;
+
+    let net: serde_json::Value = make_network(&client, "net").await.json().await;
+    let location_id = net["id"].as_i64().unwrap();
+    let _device_pubkey = create_admin_device(&client).await;
+    let posture = create_posture(&client, "Posture").await;
+    drain_gateway_events(&mut gateway_rx);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/network/{location_id}/postures"))
+        .json(&AssignPosturesData {
+            postures: vec![posture.id],
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    expect_network_modified_peers(&mut gateway_rx, location_id, &[]);
+    let events = client.drain_all_events();
+    assert!(matches!(
+        events[0].0,
+        ApiEventType::LocationPosturesAssigned { .. }
+    ));
+}
+
+#[sqlx::test]
+async fn test_removing_last_posture_refreshes_gateway_with_direct_peers(
+    _: PgPoolOptions,
+    options: PgConnectOptions,
+) {
+    let (mut client, state) = setup(options).await;
+    let mut gateway_rx = state.gateway_rx;
+
+    let net: serde_json::Value = make_network(&client, "net").await.json().await;
+    let location_id = net["id"].as_i64().unwrap();
+    let device_pubkey = create_admin_device(&client).await;
+    let posture = create_posture(&client, "Posture").await;
+    drain_gateway_events(&mut gateway_rx);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/network/{location_id}/postures"))
+        .json(&AssignPosturesData {
+            postures: vec![posture.id],
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+    expect_network_modified_peers(&mut gateway_rx, location_id, &[]);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/network/{location_id}/postures"))
+        .json(&AssignPosturesData {
+            postures: Vec::new(),
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    expect_network_modified_peers(&mut gateway_rx, location_id, &[device_pubkey]);
+    let events = client.drain_all_events();
+    assert!(matches!(
+        events[0].0,
+        ApiEventType::LocationPosturesAssigned { .. }
+    ));
+}
+
+#[sqlx::test]
+async fn test_reassigning_posture_locations_refreshes_old_and_new_locations(
+    _: PgPoolOptions,
+    options: PgConnectOptions,
+) {
+    let (mut client, state) = setup(options).await;
+    let mut gateway_rx = state.gateway_rx;
+
+    let net1: serde_json::Value = make_network(&client, "net1").await.json().await;
+    let net2: serde_json::Value = make_network(&client, "net2").await.json().await;
+    let loc1 = net1["id"].as_i64().unwrap();
+    let loc2 = net2["id"].as_i64().unwrap();
+    let device_pubkey = create_admin_device(&client).await;
+    let posture = create_posture(&client, "Posture").await;
+    drain_gateway_events(&mut gateway_rx);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/device-posture/{}/locations", posture.id))
+        .json(&AssignLocationsData {
+            locations: vec![loc1],
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+    expect_network_modified_peers(&mut gateway_rx, loc1, &[]);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/device-posture/{}/locations", posture.id))
+        .json(&AssignLocationsData {
+            locations: vec![loc2],
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    expect_network_modified_peers(&mut gateway_rx, loc1, &[device_pubkey]);
+    expect_network_modified_peers(&mut gateway_rx, loc2, &[]);
+    let events = client.drain_all_events();
+    assert!(matches!(
+        events[0].0,
+        ApiEventType::DevicePostureLocationsAssigned { .. }
+    ));
+}
+
+#[sqlx::test]
+async fn test_deleting_assigned_posture_refreshes_gateway_with_direct_peers(
+    _: PgPoolOptions,
+    options: PgConnectOptions,
+) {
+    let (mut client, state) = setup(options).await;
+    let mut gateway_rx = state.gateway_rx;
+
+    let net: serde_json::Value = make_network(&client, "net").await.json().await;
+    let location_id = net["id"].as_i64().unwrap();
+    let device_pubkey = create_admin_device(&client).await;
+    let posture = create_posture(&client, "Posture").await;
+    drain_gateway_events(&mut gateway_rx);
+    client.drain_all_events();
+
+    let response = client
+        .put(format!("/api/v1/network/{location_id}/postures"))
+        .json(&AssignPosturesData {
+            postures: vec![posture.id],
+        })
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+    expect_network_modified_peers(&mut gateway_rx, location_id, &[]);
+    client.drain_all_events();
+
+    let response = client
+        .delete(format!("/api/v1/device-posture/{}", posture.id))
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    expect_network_modified_peers(&mut gateway_rx, location_id, &[device_pubkey]);
+    let events = client.drain_all_events();
+    assert!(matches!(
+        events[0].0,
+        ApiEventType::DevicePostureDeleted { .. }
+    ));
+}
+
 #[sqlx::test]
 async fn test_device_posture_assignment_not_found(_: PgPoolOptions, options: PgConnectOptions) {
     let (mut client, _) = setup(options).await;