perf: use predefined function for context check policies (#959)

<!--
Pull requests are squashed and merged using:
- their title as the commit message
- their description as the commit body

Having a good title and description is important for the users to get
readable changelog.
-->

<!-- 1. Explain WHAT the change is about -->

Improve performance by running the functions in the main thread when it
can be done securely:
- Use predefined function for context check policies
- Use context check policy for `admin_only`

<!-- 2. Explain WHY the change cannot be made simpler -->



<!-- 3. Explain HOW users should update their code -->

#### Migration notes

---

- [ ] The change comes with new or modified tests
- [ ] Hard-to-understand functions have explanatory comments
- [ ] End-user documentation is updated to reflect the change


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
  - Enhanced predefined function handling with optional parameters.
  - Introduced a more flexible context checking mechanism.
  - Simplified policy definition for admin access.

- **Improvements**
  - Updated runtime function registration process.
  - Improved type safety for predefined function validation.
  - Streamlined error handling for function materialization.

- **Changes**
  - Removed deprecated error handling functions.
  - Modified internal representations of predefined functions.
  - Updated function signatures for predefined Deno functions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: Natoandro

src/common/src/typegraph/runtimes/deno.rs

@@ -1,6 +1,7 @@
 // Copyright Metatype OÜ, licensed under the Mozilla Public License Version 2.0.
 // SPDX-License-Identifier: MPL-2.0
 
+use anyhow::{anyhow, Context, Result};
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
@@ -11,6 +12,47 @@ pub struct FunctionMatData {
     pub script: String,
 }
 
+#[derive(Serialize, Deserialize, Clone, Debug, Hash, PartialEq, Eq)]
+#[serde(tag = "type", content = "value")]
+#[serde(rename_all = "snake_case")]
+pub enum ContextCheckX {
+    NotNull,
+    Value(String),
+    Pattern(String),
+}
+
+#[derive(PartialEq, Eq, Hash, Serialize, Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+#[serde(tag = "name", content = "param")]
+pub enum PredefinedFunctionMatData {
+    Identity,
+    True,
+    False,
+    Allow,
+    Deny,
+    Pass,
+    InternalPolicy,
+    ContextCheck { key: String, value: ContextCheckX },
+}
+
+#[derive(Serialize)]
+struct PredefinedFunctionMatDataRaw {
+    name: String,
+    param: Option<Value>,
+}
+
+impl PredefinedFunctionMatData {
+    pub fn from_raw(name: String, param: Option<String>) -> Result<Self> {
+        let param = param
+            .map(|p| serde_json::from_str(&p))
+            .transpose()
+            .context("invalid predefined function materializer parameter")?;
+        let value = serde_json::to_value(&PredefinedFunctionMatDataRaw { name, param })?;
+        serde_json::from_value(value)
+            .map_err(|e| anyhow!("invalid predefined function materializer: {e:?}"))
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct ModuleMatData {

src/typegate/src/runtimes/deno/deno.ts

@@ -30,15 +30,42 @@ import { WorkerManager } from "./worker_manager.ts";
 
 const logger = getLogger(import.meta);
 
-const predefinedFuncs: Record<string, Resolver<Record<string, unknown>>> = {
-  identity: ({ _, ...args }) => args,
-  true: () => true,
-  false: () => false,
-  allow: () => "ALLOW" as PolicyResolverOutput,
-  deny: () => "DENY" as PolicyResolverOutput,
-  pass: () => "PASS" as PolicyResolverOutput,
-  internal_policy: ({ _: { context } }) =>
+const predefinedFuncs: Record<
+  string,
+  (param: any) => Resolver<Record<string, unknown>>
+> = {
+  identity: () => ({ _, ...args }) => args,
+  true: () => () => true,
+  false: () => () => false,
+  allow: () => () => "ALLOW" as PolicyResolverOutput,
+  deny: () => () => "DENY" as PolicyResolverOutput,
+  pass: () => () => "PASS" as PolicyResolverOutput,
+  internal_policy: () => ({ _: { context } }) =>
     context.provider === "internal" ? "ALLOW" : "PASS" as PolicyResolverOutput,
+  context_check: ({ key, value }) => {
+    let check: (value: any) => boolean;
+    switch (value.type) {
+      case "not_null":
+        check = (v) => v != null;
+        break;
+      case "value":
+        check = (v) => v === value.value;
+        break;
+      case "pattern":
+        check = (v) => v != null && new RegExp(value.value).test(v);
+        break;
+      default:
+        throw new Error("unreachable");
+    }
+    const path = key.split(".");
+    return ({ _: { context } }) => {
+      let value: any = context;
+      for (const segment of path) {
+        value = value?.[segment];
+      }
+      return check(value) ? "PASS" : "DENY" as PolicyResolverOutput;
+    };
+  },
 };
 
 export class DenoRuntime extends Runtime {
@@ -229,7 +256,7 @@ export class DenoRuntime extends Runtime {
       if (!func) {
         throw new Error(`predefined function ${mat.data.name} not found`);
       }
-      return func;
+      return func(mat.data.param);
     }
 
     if (mat.name === "static") {

src/typegate/src/typegraphs/typegate.json

@@ -807,15 +807,21 @@
       "data": {}
     },
     {
-      "name": "function",
+      "name": "predefined_function",
       "runtime": 0,
       "effect": {
         "effect": "read",
         "idempotent": true
       },
       "data": {
-        "script": "var _my_lambda = (_args, { context }) => context.username === 'admin' ? 'ALLOW' : 'DENY' ",
-        "secrets": []
+        "name": "context_check",
+        "param": {
+          "key": "username",
+          "value": {
+            "type": "value",
+            "value": "admin"
+          }
+        }
       }
     },
     {
@@ -942,7 +948,7 @@
   ],
   "policies": [
     {
-      "name": "admin_only",
+      "name": "__ctx_username_admin",
       "materializer": 1
     }
   ],

src/typegate/src/typegraphs/typegate.py

@@ -1,12 +1,11 @@
 # Copyright Metatype OÜ, licensed under the Mozilla Public License Version 2.0.
 # SPDX-License-Identifier: MPL-2.0
 
-from typegraph import Graph, fx, t, typegraph
+from typegraph import Graph, fx, t, typegraph, Policy
 from typegraph.gen.exports.runtimes import TypegateOperation
 from typegraph.gen.types import Err
 from typegraph.graph.params import Auth, Cors, Rate
 from typegraph.runtimes.base import Materializer
-from typegraph.runtimes.deno import DenoRuntime
 from typegraph.wit import runtimes, store
 
 ### Prisma query (Json protocol):
@@ -95,11 +94,7 @@
     ),
 )
 def typegate(g: Graph):
-    deno = DenoRuntime()
-    admin_only = deno.policy(
-        "admin_only",
-        code="(_args, { context }) => context.username === 'admin' ? 'ALLOW' : 'DENY' ",
-    )
+    admin_only = Policy.context("username", "admin")
 
     g.auth(Auth.basic(["admin"]))
 

src/typegraph/core/src/conversion/runtimes.rs

@@ -128,10 +128,7 @@ impl MaterializerConverter for DenoMaterializer {
                 ("import_function".to_string(), data)
             }
             Predefined(predef) => {
-                let data = serde_json::from_value(json!({
-                    "name": predef.name,
-                }))
-                .unwrap();
+                let data = serde_json::from_value(serde_json::to_value(predef).unwrap()).unwrap();
                 ("predefined_function".to_string(), data)
             }
         };

src/typegraph/core/src/errors.rs

@@ -143,10 +143,6 @@ pub fn object_not_found(kind: &str, id: u32) -> TgError {
 //     .into()
 // }
 
-pub fn unknown_predefined_function(name: &str, runtime: &str) -> TgError {
-    format!("unknown predefined function {name} for runtime {runtime}").into()
-}
-
 pub fn duplicate_policy_name(name: &str) -> TgError {
     format!("duplicate policy name '{name}'").into()
 }

src/typegraph/core/src/global_store.rs

@@ -14,6 +14,7 @@ use crate::wit::utils::Auth as WitAuth;
 
 #[allow(unused)]
 use crate::wit::runtimes::{Effect, MaterializerDenoPredefined, MaterializerId};
+use common::typegraph::runtimes::deno::PredefinedFunctionMatData;
 use graphql_parser::parse_query;
 use indexmap::IndexMap;
 use std::rc::Rc;
@@ -55,7 +56,7 @@ pub struct Store {
     pub policies: Vec<Policy>,
 
     deno_runtime: RuntimeId,
-    predefined_deno_functions: HashMap<String, MaterializerId>,
+    predefined_deno_functions: HashMap<PredefinedFunctionMatData, MaterializerId>,
     deno_modules: IndexMap<String, MaterializerId>,
 
     public_policy_id: PolicyId,
@@ -88,9 +89,7 @@ impl Store {
                 runtime_id: deno_runtime,
                 effect: Effect::Read,
                 data: MaterializerData::Deno(Rc::new(DenoMaterializer::Predefined(
-                    crate::wit::runtimes::MaterializerDenoPredefined {
-                        name: "pass".to_string(),
-                    },
+                    PredefinedFunctionMatData::Pass,
                 ))),
             }],
 
@@ -104,8 +103,6 @@ impl Store {
     }
 }
 
-const PREDEFINED_DENO_FUNCTIONS: &[&str] = &["identity", "true"];
-
 thread_local! {
     pub static STORE: RefCell<Store> = RefCell::new(Store::new());
     pub static SDK_VERSION: String = "0.5.0-rc.9".to_owned();
@@ -350,25 +347,24 @@ impl Store {
         with_store(|s| s.public_policy_id)
     }
 
-    pub fn get_predefined_deno_function(name: String) -> Result<MaterializerId> {
-        if let Some(mat) = with_store(|s| s.predefined_deno_functions.get(&name).cloned()) {
-            Ok(mat)
-        } else if !PREDEFINED_DENO_FUNCTIONS.iter().any(|n| n == &name) {
-            Err(errors::unknown_predefined_function(&name, "deno"))
+    pub fn get_predefined_deno_function(
+        name: String,
+        param: Option<String>,
+    ) -> Result<MaterializerId> {
+        let mat = PredefinedFunctionMatData::from_raw(name, param)?;
+        if let Some(mat_id) = with_store(|s| s.predefined_deno_functions.get(&mat).cloned()) {
+            Ok(mat_id)
         } else {
             let runtime_id = Store::get_deno_runtime();
-            let mat = Store::register_materializer(Materializer {
+            let mat_id = Store::register_materializer(Materializer {
                 runtime_id,
                 effect: Effect::Read,
-                data: Rc::new(DenoMaterializer::Predefined(MaterializerDenoPredefined {
-                    name: name.clone(),
-                }))
-                .into(),
+                data: Rc::new(DenoMaterializer::Predefined(mat.clone())).into(),
             });
             with_store_mut(|s| {
-                s.predefined_deno_functions.insert(name, mat);
+                s.predefined_deno_functions.insert(mat, mat_id);
             });
-            Ok(mat)
+            Ok(mat_id)
         }
     }
 

src/typegraph/core/src/lib.rs

@@ -19,10 +19,10 @@ mod test_utils;
 
 use std::collections::HashSet;
 
+use common::typegraph::runtimes::deno::{ContextCheckX, PredefinedFunctionMatData};
 use common::typegraph::Injection;
 use errors::{Result, TgError};
 use global_store::Store;
-use indoc::formatdoc;
 use params::apply;
 use regex::Regex;
 use runtimes::{DenoMaterializer, Materializer};
@@ -38,7 +38,7 @@ use wit::core::{
     TypeEither, TypeFile, TypeFloat, TypeFunc, TypeId as CoreTypeId, TypeInteger, TypeList,
     TypeOptional, TypeString, TypeStruct, TypeUnion, TypegraphInitParams,
 };
-use wit::runtimes::{Guest, MaterializerDenoFunc};
+use wit::runtimes::{Guest as _, MaterializerDenoPredefined};
 
 pub mod wit {
     wit_bindgen::generate!({
@@ -52,6 +52,17 @@ pub mod wit {
 
 pub struct Lib {}
 
+impl From<ContextCheck> for ContextCheckX {
+    fn from(check: ContextCheck) -> Self {
+        use ContextCheck as CC;
+        match check {
+            CC::NotNull => ContextCheckX::NotNull,
+            CC::Value(v) => ContextCheckX::Value(v),
+            CC::Pattern(p) => ContextCheckX::Pattern(p),
+        }
+    }
+}
+
 impl wit::core::Guest for Lib {
     fn init_typegraph(params: TypegraphInitParams) -> Result<()> {
         typegraph::init(params)
@@ -185,9 +196,7 @@ impl wit::core::Guest for Lib {
     }
 
     fn get_internal_policy() -> Result<(PolicyId, String)> {
-        let deno_mat = DenoMaterializer::Predefined(wit::runtimes::MaterializerDenoPredefined {
-            name: "internal_policy".to_string(),
-        });
+        let deno_mat = DenoMaterializer::Predefined(PredefinedFunctionMatData::InternalPolicy);
         let mat = Materializer::deno(deno_mat, crate::wit::runtimes::Effect::Read);
         let policy_id = Store::register_policy(
             Policy {
@@ -213,42 +222,19 @@ impl wit::core::Guest for Lib {
             .replace_all(&name, "_")
             .to_string();
 
-        let check = match check {
-            ContextCheck::NotNull => "value != null ? 'PASS' : 'DENY'".to_string(),
-            ContextCheck::Value(val) => {
-                format!(
-                    "value === {} ?  'PASS' : 'DENY'",
-                    serde_json::to_string(&val).unwrap()
-                )
-            }
-            ContextCheck::Pattern(pattern) => {
-                format!(
-                    "new RegExp({}).test(value) ? 'PASS' : 'DENY' ",
-                    serde_json::to_string(&pattern).unwrap()
-                )
-            }
-        };
+        let check: ContextCheckX = check.into();
+        let check = serde_json::json!({
+            "key": key,
+            "value": check
+        });
 
-        let key = serde_json::to_string(&key).unwrap();
-
-        let code = formatdoc! {r#"
-            (_, {{ context }}) => {{
-                const chunks = {key}.split(".");
-                let value = context;
-                for (const chunk of chunks) {{
-                    value = value?.[chunk];
-                }}
-                return {check};
-            }}
-        "# };
-
-        let mat_id = Lib::register_deno_func(
-            MaterializerDenoFunc {
-                code,
-                secrets: vec![],
-            },
-            wit::runtimes::Effect::Read,
-        )?;
+        let mat_id = Lib::get_predefined_deno_func(MaterializerDenoPredefined {
+            name: "context_check".to_string(),
+            param: Some(
+                serde_json::to_string(&check)
+                    .map_err(|e| format!("Error while serializing context check: {e:?}"))?,
+            ),
+        })?;
 
         Lib::register_policy(Policy {
             name: name.clone(),

src/typegraph/core/src/runtimes/deno.rs

@@ -1,6 +1,8 @@
 // Copyright Metatype OÜ, licensed under the Mozilla Public License Version 2.0.
 // SPDX-License-Identifier: MPL-2.0
 
+use common::typegraph::runtimes::deno::PredefinedFunctionMatData;
+
 use crate::wit::runtimes as wit;
 
 #[derive(Debug)]
@@ -25,7 +27,7 @@ pub struct MaterializerDenoStatic {
 pub enum DenoMaterializer {
     Static(MaterializerDenoStatic),
     Inline(wit::MaterializerDenoFunc),
-    Predefined(wit::MaterializerDenoPredefined),
+    Predefined(PredefinedFunctionMatData),
     Module(MaterializerDenoModule),
     Import(MaterializerDenoImport),
 }