Embed the default policy in the binary

Dockerfile

@@ -30,17 +30,27 @@ RUN find public -type f -exec touch -t 197001010000.00 {} +
 ## Build stage that builds the OPA policies ##
 FROM --platform=${BUILDPLATFORM} docker.io/library/debian:${DEBIAN_VERSION_NAME}-slim AS policy
 
+# Install make
+RUN rm -f /etc/apt/apt.conf.d/docker-clean; echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache
+RUN \
+  --mount=type=cache,sharing=locked,target=/var/cache/apt \
+  --mount=type=cache,sharing=locked,target=/var/lib/apt \
+  apt update && apt install -y --no-install-recommends \
+  make
+
 ARG BUILDOS
 ARG BUILDARCH
 ARG OPA_VERSION
 
+# Download Open Policy Agent
 ADD --chmod=755 https://github.com/open-policy-agent/opa/releases/download/v${OPA_VERSION}/opa_${BUILDOS}_${BUILDARCH}_static /usr/local/bin/opa
 
-WORKDIR /policies
-COPY ./policies/ /policies
-RUN opa build -t wasm -e "client_registration/allow" -e "login/allow" -e "register/allow" client_registration.rego login.rego register.rego \
-  && tar xzf bundle.tar.gz /policy.wasm \
-  && rm -f bundle.tar.gz
+WORKDIR /app/crates/policy/policies
+COPY ./crates/policy/policies/ /app/crates/policy/policies
+RUN make -B
+
+# Change the timestamp of built files for better caching
+RUN touch -t 197001010000.00 {} policy.wasm
 
 ## Base image with cargo-chef and the right cross-compilation toolchain ##
 # cargo-chef helps with caching dependencies between builds
@@ -108,6 +118,7 @@ RUN \
 COPY ./Cargo.toml ./Cargo.lock /app/
 COPY ./crates /app/crates
 COPY --from=static-files /app/crates/static-files/public /app/crates/static-files/public
+COPY --from=policy /app/crates/policy/policies/policy.wasm /app/crates/policy/policies/policy.wasm
 ENV SQLX_OFFLINE=true
 RUN \
   --mount=type=cache,sharing=private,target=/usr/local/cargo/registry \
@@ -122,13 +133,11 @@ RUN mv target/$(/docker-arch-to-rust-target.sh "${TARGETPLATFORM}")/release/mas-
 ## Runtime stage, debug variant ##
 FROM --platform=${TARGETPLATFORM} gcr.io/distroless/cc-debian${DEBIAN_VERSION}:debug-nonroot AS debug
 COPY --from=builder /usr/local/bin/mas-cli /usr/local/bin/mas-cli
-COPY --chmod=444 --from=policy /policies/policy.wasm /policies/policy.wasm
 WORKDIR /
 ENTRYPOINT ["/mas-cli"]
 
 ## Runtime stage ##
 FROM --platform=${TARGETPLATFORM} gcr.io/distroless/cc-debian${DEBIAN_VERSION}:nonroot
 COPY --from=builder /usr/local/bin/mas-cli /usr/local/bin/mas-cli
-COPY --chmod=444 --from=policy /policies/policy.wasm /policies/policy.wasm
 WORKDIR /
 ENTRYPOINT ["/usr/local/bin/mas-cli"]

README.md

@@ -9,6 +9,7 @@ See the [Documentation](https://matrix-org.github.io/matrix-authentication-servi
 
 - [Install Rust and Cargo](https://www.rust-lang.org/learn/get-started)
 - [Install Node.js and npm](https://nodejs.org/)
+- [Install Open Policy Agent](https://www.openpolicyagent.org/docs/latest/#1-download-opa)
 - Clone this repository
 - Generate the frontend:
   ```sh
@@ -17,6 +18,14 @@ See the [Documentation](https://matrix-org.github.io/matrix-authentication-servi
   npm run build
   cd ../..
   ```
+- Build the Open Policy Agent policies
+  ```sh
+  cd crates/policy/policies
+  make
+  # OR, if you don't have `opa` installed and want to build through the OPA docker image
+  make DOCKER=1
+  cd ../../..
+  ```
 - Generate the sample config via `cargo run -- config generate > config.yaml`
 - Run the database migrations via `cargo run -- database migrate`
 - Run the server via `cargo run -- server -c config.yaml`

crates/cli/src/commands/server.rs

@@ -30,6 +30,7 @@ use mas_router::UrlBuilder;
 use mas_storage::MIGRATOR;
 use mas_tasks::TaskQueue;
 use mas_templates::Templates;
+use tokio::io::AsyncRead;
 use tracing::{error, info};
 
 #[derive(Parser, Debug, Default)]
@@ -177,10 +178,19 @@ impl Options {
 
         let encrypter = config.secrets.encrypter();
 
-        // Load and compile the WASM policies
-        let mut policy = tokio::fs::File::open(&config.policy.wasm_module)
-            .await
-            .context("failed to open OPA WASM policy file")?;
+        // Load and compile the WASM policies (and fallback to the default embedded one)
+        info!("Loading and compiling the policy module");
+        let mut policy: Box<dyn AsyncRead + std::marker::Unpin> =
+            if let Some(path) = &config.policy.wasm_module {
+                Box::new(
+                    tokio::fs::File::open(path)
+                        .await
+                        .context("failed to open OPA WASM policy file")?,
+                )
+            } else {
+                Box::new(mas_policy::default_wasm_policy())
+            };
+
         let policy_factory = PolicyFactory::load(
             &mut policy,
             config.policy.data.clone().unwrap_or_default(),
@@ -188,7 +198,8 @@ impl Options {
             config.policy.register_entrypoint.clone(),
             config.policy.client_registration_entrypoint.clone(),
         )
-        .await?;
+        .await
+        .context("failed to load the policy")?;
         let policy_factory = Arc::new(policy_factory);
 
         // Load and compile the templates

crates/config/src/sections/policy.rs

@@ -21,10 +21,6 @@ use serde_with::serde_as;
 
 use super::ConfigurationSection;
 
-fn default_wasm_module() -> PathBuf {
-    "./policies/policy.wasm".into()
-}
-
 fn default_client_registration_endpoint() -> String {
     "client_registration/allow".to_string()
 }
@@ -42,8 +38,8 @@ fn default_register_endpoint() -> String {
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
 pub struct PolicyConfig {
     /// Path to the WASM module
-    #[serde(default = "default_wasm_module")]
-    pub wasm_module: PathBuf,
+    #[serde(default)]
+    pub wasm_module: Option<PathBuf>,
 
     /// Entrypoint to use when evaluating client registrations
     #[serde(default = "default_client_registration_endpoint")]
@@ -65,7 +61,7 @@ pub struct PolicyConfig {
 impl Default for PolicyConfig {
     fn default() -> Self {
         Self {
-            wasm_module: default_wasm_module(),
+            wasm_module: None,
             client_registration_entrypoint: default_client_registration_endpoint(),
             login_entrypoint: default_login_endpoint(),
             register_entrypoint: default_register_endpoint(),

crates/handlers/src/oauth2/registration.rs

@@ -118,7 +118,7 @@ pub(crate) async fn post(
         return Err(RouteError::InvalidClientMetadata);
     }
 
-    let mut policy = policy_factory.instanciate().await?;
+    let mut policy = policy_factory.instantiate().await?;
     let allowed = policy.evaluate_client_registration(&body).await?;
     if !allowed {
         return Err(RouteError::PolicyDenied);

crates/handlers/src/views/register.rs

@@ -131,7 +131,7 @@ pub(crate) async fn post(
             state.add_error_on_field(RegisterFormField::PasswordConfirm, FieldError::Unspecified);
         }
 
-        let mut policy = policy_factory.instanciate().await?;
+        let mut policy = policy_factory.instantiate().await?;
         let res = policy
             .evaluate_register(&form.username, &form.email)
             .await?;

crates/policy/policies/.gitignore

@@ -0,0 +1,2 @@
+/policy.wasm
+/bundle.tar.gz

crates/policy/policies/Makefile

@@ -0,0 +1,27 @@
+# Set to 1 to run OPA through Docker
+DOCKER := 0
+
+ifeq ($(DOCKER), 0)
+	OPA := opa
+else
+	OPA := docker run -v $(shell pwd):/policies -w /policies --rm docker.io/openpolicyagent/opa:0.40.0
+endif
+
+policy.wasm: client_registration.rego login.rego register.rego
+	$(OPA) build -t wasm -e "client_registration/allow" -e "login/allow" -e "register/allow" $^
+	tar xzf bundle.tar.gz /policy.wasm
+	$(RM) bundle.tar.gz
+	touch $@
+
+.PHONY: fmt
+fmt:
+	$(OPA) fmt -w .
+
+.PHONY: test
+test:
+	$(OPA) test -v .
+
+.PHONY: lint
+lint:
+	$(OPA) fmt -d --fail .
+	$(OPA) check --strict .

crates/policy/src/lib.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::io::Cursor;
+
 use anyhow::bail;
 use oauth2_types::registration::ClientMetadata;
 use opa_wasm::Runtime;
@@ -20,6 +22,12 @@ use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncReadExt};
 use wasmtime::{Config, Engine, Module, Store};
 
+const DEFAULT_POLICY: &[u8] = include_bytes!("../policies/policy.wasm");
+
+pub fn default_wasm_policy() -> impl AsyncRead + std::marker::Unpin {
+    Cursor::new(DEFAULT_POLICY)
+}
+
 #[derive(Debug, Error)]
 pub enum LoadError {
     #[error("failed to read module")]
@@ -33,6 +41,9 @@ pub enum LoadError {
 
     #[error("failed to compile WASM module")]
     Compilation(#[source] anyhow::Error),
+
+    #[error("failed to instantiate a test instance")]
+    Instantiate(#[source] anyhow::Error),
 }
 
 pub struct PolicyFactory {
@@ -55,27 +66,39 @@ impl PolicyFactory {
         let mut config = Config::default();
         config.async_support(true);
         config.cranelift_opt_level(wasmtime::OptLevel::Speed);
+
         let engine = Engine::new(&config).map_err(LoadError::Engine)?;
+
+        // Read and compile the module
         let mut buf = Vec::new();
         source.read_to_end(&mut buf).await?;
+        // Compilation is CPU-bound, so spawn that in a blocking task
         let (engine, module) = tokio::task::spawn_blocking(move || {
-            let module = Module::new(&engine, buf);
-            (engine, module)
+            let module = Module::new(&engine, buf)?;
+            anyhow::Ok((engine, module))
         })
-        .await?;
-        let module = module.map_err(LoadError::Compilation)?;
+        .await?
+        .map_err(LoadError::Compilation)?;
 
-        Ok(Self {
+        let factory = Self {
             engine,
             module,
             data,
             login_entrypoint,
             register_entrypoint,
             client_registration_entrypoint,
-        })
+        };
+
+        // Try to instanciate
+        factory
+            .instantiate()
+            .await
+            .map_err(LoadError::Instantiate)?;
+
+        Ok(factory)
     }
 
-    pub async fn instanciate(&self) -> Result<Policy, anyhow::Error> {
+    pub async fn instantiate(&self) -> Result<Policy, anyhow::Error> {
         let mut store = Store::new(&self.engine, ());
         let runtime = Runtime::new(&mut store, &self.module).await?;
 

policies/Makefile

@@ -1,13 +0,0 @@
-policy.wasm: client_registration.rego login.rego register.rego
-	opa build -t wasm -e "client_registration/allow" -e "login/allow" -e "register/allow" $^
-	tar xzf bundle.tar.gz /policy.wasm
-	rm -f bundle.tar.gz
-	touch $@
-
-.PHONY: fmt
-fmt:
-	opa fmt -w .
-
-.PHONY: test
-test:
-	opa test -v .