Admin API to dynamically set policy data (#4115)

Author: sandhose

Cargo.lock

@@ -61,6 +61,7 @@ syn2mas = { path = "./crates/syn2mas", version = "=0.14.1" }
 version = "0.14.1"
 features = ["axum", "axum-extra", "axum-json", "axum-query", "macros"]
 
+# An `Arc` that can be atomically updated
 [workspace.dependencies.arc-swap]
 version = "1.7.1"
 

Cargo.toml

@@ -203,6 +203,12 @@ impl FromRef<AppState> for Limiter {
     }
 }
 
+impl FromRef<AppState> for Arc<PolicyFactory> {
+    fn from_ref(input: &AppState) -> Self {
+        input.policy_factory.clone()
+    }
+}
+
 impl FromRef<AppState> for Arc<dyn HomeserverConnection> {
     fn from_ref(input: &AppState) -> Self {
         Arc::clone(&input.homeserver_connection)

crates/cli/src/app_state.rs

@@ -1,4 +1,4 @@
-// Copyright 2024 New Vector Ltd.
+// Copyright 2024, 2025 New Vector Ltd.
 // Copyright 2022-2024 The Matrix.org Foundation C.I.C.
 //
 // SPDX-License-Identifier: AGPL-3.0-only
@@ -8,10 +8,14 @@ use std::process::ExitCode;
 
 use clap::Parser;
 use figment::Figment;
-use mas_config::{ConfigurationSection, ConfigurationSectionExt, MatrixConfig, PolicyConfig};
+use mas_config::{
+    ConfigurationSection, ConfigurationSectionExt, DatabaseConfig, MatrixConfig, PolicyConfig,
+};
 use tracing::{info, info_span};
 
-use crate::util::policy_factory_from_config;
+use crate::util::{
+    database_pool_from_config, load_policy_factory_dynamic_data, policy_factory_from_config,
+};
 
 #[derive(Parser, Debug)]
 pub(super) struct Options {
@@ -22,21 +26,31 @@ pub(super) struct Options {
 #[derive(Parser, Debug)]
 enum Subcommand {
     /// Check that the policies compile
-    Policy,
+    Policy {
+        /// With dynamic data loaded
+        #[arg(long)]
+        with_dynamic_data: bool,
+    },
 }
 
 impl Options {
     #[tracing::instrument(skip_all)]
     pub async fn run(self, figment: &Figment) -> anyhow::Result<ExitCode> {
         use Subcommand as SC;
         match self.subcommand {
-            SC::Policy => {
+            SC::Policy { with_dynamic_data } => {
                 let _span = info_span!("cli.debug.policy").entered();
                 let config = PolicyConfig::extract_or_default(figment)?;
                 let matrix_config = MatrixConfig::extract(figment)?;
                 info!("Loading and compiling the policy module");
                 let policy_factory = policy_factory_from_config(&config, &matrix_config).await?;
 
+                if with_dynamic_data {
+                    let database_config = DatabaseConfig::extract(figment)?;
+                    let pool = database_pool_from_config(&database_config).await?;
+                    load_policy_factory_dynamic_data(&policy_factory, &pool).await?;
+                }
+
                 let _instance = policy_factory.instantiate().await?;
             }
         }

crates/cli/src/commands/debug.rs

@@ -25,7 +25,8 @@ use crate::{
     app_state::AppState,
     lifecycle::LifecycleManager,
     util::{
-        database_pool_from_config, homeserver_connection_from_config, mailer_from_config,
+        database_pool_from_config, homeserver_connection_from_config,
+        load_policy_factory_dynamic_data_continuously, mailer_from_config,
         password_manager_from_config, policy_factory_from_config, site_config_from_config,
         templates_from_config, test_mailer_in_background,
     },
@@ -129,6 +130,14 @@ impl Options {
         let policy_factory = policy_factory_from_config(&config.policy, &config.matrix).await?;
         let policy_factory = Arc::new(policy_factory);
 
+        load_policy_factory_dynamic_data_continuously(
+            &policy_factory,
+            &pool,
+            shutdown.soft_shutdown_token(),
+            shutdown.task_tracker(),
+        )
+        .await?;
+
         let url_builder = UrlBuilder::new(
             config.http.public_base.clone(),
             config.http.issuer.clone(),

crates/cli/src/commands/server.rs

@@ -19,11 +19,14 @@ use mas_matrix::{HomeserverConnection, ReadOnlyHomeserverConnection};
 use mas_matrix_synapse::SynapseConnection;
 use mas_policy::PolicyFactory;
 use mas_router::UrlBuilder;
+use mas_storage::RepositoryAccess;
+use mas_storage_pg::PgRepository;
 use mas_templates::{SiteConfigExt, TemplateLoadingError, Templates};
 use sqlx::{
     ConnectOptions, PgConnection, PgPool,
     postgres::{PgConnectOptions, PgPoolOptions},
 };
+use tokio_util::{sync::CancellationToken, task::TaskTracker};
 use tracing::{Instrument, log::LevelFilter};
 
 pub async fn password_manager_from_config(
@@ -377,6 +380,66 @@ pub async fn database_connection_from_config_with_options(
         .context("could not connect to the database")
 }
 
+/// Update the policy factory dynamic data from the database and spawn a task to
+/// periodically update it
+// XXX: this could be put somewhere else?
+pub async fn load_policy_factory_dynamic_data_continuously(
+    policy_factory: &Arc<PolicyFactory>,
+    pool: &PgPool,
+    cancellation_token: CancellationToken,
+    task_tracker: &TaskTracker,
+) -> Result<(), anyhow::Error> {
+    let policy_factory = policy_factory.clone();
+    let pool = pool.clone();
+
+    load_policy_factory_dynamic_data(&policy_factory, &pool).await?;
+
+    task_tracker.spawn(async move {
+        let mut interval = tokio::time::interval(Duration::from_secs(60));
+
+        loop {
+            tokio::select! {
+                () = cancellation_token.cancelled() => {
+                    return;
+                }
+                _ = interval.tick() => {}
+            }
+
+            if let Err(err) = load_policy_factory_dynamic_data(&policy_factory, &pool).await {
+                tracing::error!(
+                    error = ?err,
+                    "Failed to load policy factory dynamic data"
+                );
+                cancellation_token.cancel();
+                return;
+            }
+        }
+    });
+
+    Ok(())
+}
+
+/// Update the policy factory dynamic data from the database
+#[tracing::instrument(name = "policy.load_dynamic_data", skip_all, err(Debug))]
+pub async fn load_policy_factory_dynamic_data(
+    policy_factory: &PolicyFactory,
+    pool: &PgPool,
+) -> Result<(), anyhow::Error> {
+    let mut repo = PgRepository::from_pool(pool)
+        .await
+        .context("Failed to acquire database connection")?;
+
+    if let Some(data) = repo.policy_data().get().await? {
+        let id = data.id;
+        let updated = policy_factory.set_dynamic_data(data).await?;
+        if updated {
+            tracing::info!(policy_data.id = %id, "Loaded dynamic policy data from the database");
+        }
+    }
+
+    Ok(())
+}
+
 /// Create a clonable, type-erased [`HomeserverConnection`] from the
 /// configuration
 pub fn homeserver_connection_from_config(

crates/cli/src/util.rs

@@ -10,6 +10,7 @@ use thiserror::Error;
 
 pub(crate) mod compat;
 pub mod oauth2;
+pub(crate) mod policy_data;
 mod site_config;
 pub(crate) mod tokens;
 pub(crate) mod upstream_oauth2;
@@ -32,6 +33,7 @@ pub use self::{
         AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, Client, DeviceCodeGrant,
         DeviceCodeGrantState, InvalidRedirectUriError, JwksOrJwksUri, Pkce, Session, SessionState,
     },
+    policy_data::PolicyData,
     site_config::{CaptchaConfig, CaptchaService, SessionExpirationConfig, SiteConfig},
     tokens::{
         AccessToken, AccessTokenState, RefreshToken, RefreshTokenState, TokenFormatError, TokenType,

crates/data-model/src/lib.rs

@@ -0,0 +1,15 @@
+// Copyright 2025 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+
+use chrono::{DateTime, Utc};
+use serde::Serialize;
+use ulid::Ulid;
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct PolicyData {
+    pub id: Ulid,
+    pub created_at: DateTime<Utc>,
+    pub data: serde_json::Value,
+}

crates/data-model/src/policy_data.rs

@@ -22,6 +22,7 @@ use indexmap::IndexMap;
 use mas_axum_utils::FancyError;
 use mas_http::CorsLayerExt;
 use mas_matrix::HomeserverConnection;
+use mas_policy::PolicyFactory;
 use mas_router::{
     ApiDoc, ApiDocCallback, OAuth2AuthorizationEndpoint, OAuth2TokenEndpoint, Route, SimpleRoute,
     UrlBuilder,
@@ -47,6 +48,11 @@ fn finish(t: TransformOpenApi) -> TransformOpenApi {
             description: Some("Manage compatibility sessions from legacy clients".to_owned()),
             ..Tag::default()
         })
+        .tag(Tag {
+            name: "policy-data".to_owned(),
+            description: Some("Manage the dynamic policy data".to_owned()),
+            ..Tag::default()
+        })
         .tag(Tag {
             name: "oauth2-session".to_owned(),
             description: Some("Manage OAuth2 sessions".to_owned()),
@@ -115,6 +121,7 @@ where
     CallContext: FromRequestParts<S>,
     Templates: FromRef<S>,
     UrlBuilder: FromRef<S>,
+    Arc<PolicyFactory>: FromRef<S>,
 {
     // We *always* want to explicitly set the possible responses, beacuse the
     // infered ones are not necessarily correct

crates/handlers/src/admin/mod.rs

@@ -534,3 +534,50 @@ impl UpstreamOAuthLink {
         ]
     }
 }
+
+/// The policy data
+#[derive(Serialize, JsonSchema)]
+pub struct PolicyData {
+    #[serde(skip)]
+    id: Ulid,
+
+    /// The creation date of the policy data
+    created_at: DateTime<Utc>,
+
+    /// The policy data content
+    data: serde_json::Value,
+}
+
+impl From<mas_data_model::PolicyData> for PolicyData {
+    fn from(policy_data: mas_data_model::PolicyData) -> Self {
+        Self {
+            id: policy_data.id,
+            created_at: policy_data.created_at,
+            data: policy_data.data,
+        }
+    }
+}
+
+impl Resource for PolicyData {
+    const KIND: &'static str = "policy-data";
+    const PATH: &'static str = "/api/admin/v1/policy-data";
+
+    fn id(&self) -> Ulid {
+        self.id
+    }
+}
+
+impl PolicyData {
+    /// Samples of policy data
+    pub fn samples() -> [Self; 1] {
+        [Self {
+            id: Ulid::from_bytes([0x01; 16]),
+            created_at: DateTime::default(),
+            data: serde_json::json!({
+                "hello": "world",
+                "foo": 42,
+                "bar": true
+            }),
+        }]
+    }
+}