Add AuthCtx to ReducerContext for rust

crates/bindings-sys/src/lib.rs

@@ -588,6 +588,7 @@ pub mod raw {
         ///
         /// - `out_ptr` is NULL or `out` is not in bounds of WASM memory.
         pub fn identity(out_ptr: *mut u8);
+
     }
 
     // See comment on previous `extern "C"` block re: ABI version.
@@ -617,6 +618,11 @@ pub mod raw {
         ///
         /// If this function returns an error, `out` is not written.
         pub fn bytes_source_remaining_length(source: BytesSource, out: *mut u32) -> i16;
+
+        /// Find the jwt payload for the given connection id, and write the
+        /// BytesSourceId to the given pointer.
+        /// If this is not found, BytesSourceId::INVALID (aka 0) will be written.
+        pub fn get_jwt(connection_id_ptr: *const u8, bytes_source_id: *mut BytesSource);
     }
 
     /// What strategy does the database index use?
@@ -1118,6 +1124,34 @@ pub fn identity() -> [u8; 32] {
     buf
 }
 
+#[inline]
+pub fn get_jwt(connection_id: [u8; 16]) -> Option<String> {
+    let mut source: raw::BytesSource = raw::BytesSource::INVALID;
+    unsafe {
+        raw::get_jwt(connection_id.as_ptr(), &mut source);
+    }
+    if source == raw::BytesSource::INVALID {
+        return None; // No JWT found.
+    }
+    let len = {
+        let mut len = 0;
+        let ret = unsafe { raw::bytes_source_remaining_length(source, &raw mut len) };
+        match ret {
+            0 => len,
+            _ => panic!("invalid source"),
+        }
+    };
+    let mut buf = vec![0u8; len as usize];
+    let mut bytes_read = len as usize;
+    let ret = unsafe { raw::bytes_source_read(source, buf.as_mut_ptr(), &mut bytes_read) };
+    // We should have exhausted the source.
+    assert_eq!(ret, -1);
+    // We should get exactly `len` bytes.
+    assert_eq!(bytes_read, len as usize);
+
+    Some(std::str::from_utf8(&buf[..bytes_read]).unwrap().to_string())
+}
+
 pub struct RowIter {
     raw: raw::RowIter,
 }

crates/bindings/Cargo.toml

@@ -35,6 +35,8 @@ rand08 = { workspace = true, optional = true }
 # if someone tries to use rand's ThreadRng, it will fail to link
 # because no one defined __getrandom_custom
 getrandom02 = { workspace = true, optional = true, features = ["custom"] }
+serde_json.workspace = true
+once_cell = "1.21.3"
 
 [dev-dependencies]
 insta.workspace = true

crates/bindings/src/lib.rs

@@ -12,12 +12,13 @@ pub mod rt;
 #[doc(hidden)]
 pub mod table;
 
-use spacetimedb_lib::bsatn;
-use std::cell::RefCell;
-
 pub use log;
 #[cfg(feature = "rand")]
 pub use rand08 as rand;
+use spacetimedb_lib::bsatn;
+use std::cell::{OnceCell, RefCell};
+use std::ops::Deref;
+use std::sync::LazyLock;
 
 #[cfg(feature = "unstable")]
 pub use client_visibility_filter::Filter;
@@ -732,6 +733,8 @@ pub struct ReducerContext {
     /// See the [`#[table]`](macro@crate::table) macro for more information.
     pub db: Local,
 
+    sender_auth: AuthCtx,
+
     #[cfg(feature = "rand08")]
     rng: std::cell::OnceCell<StdbRng>,
 }
@@ -744,10 +747,32 @@ impl ReducerContext {
             sender: Identity::__dummy(),
             timestamp: Timestamp::UNIX_EPOCH,
             connection_id: None,
+            sender_auth: AuthCtx::internal(),
             rng: std::cell::OnceCell::new(),
         }
     }
 
+    #[doc(hidden)]
+    fn new(db: Local, sender: Identity, connection_id: Option<ConnectionId>, timestamp: Timestamp) -> Self {
+        let sender_auth = match connection_id {
+            Some(cid) => AuthCtx::from_connection_id(cid),
+            None => AuthCtx::internal(),
+        };
+        Self {
+            db,
+            sender,
+            timestamp,
+            connection_id,
+            sender_auth,
+            #[cfg(feature = "rand08")]
+            rng: std::cell::OnceCell::new(),
+        }
+    }
+
+    pub fn sender_auth(&self) -> &AuthCtx {
+        &self.sender_auth
+    }
+
     /// Read the current module's [`Identity`].
     pub fn identity(&self) -> Identity {
         // Hypothetically, we *could* read the module identity out of the system tables.
@@ -796,6 +821,114 @@ impl DbContext for ReducerContext {
 #[non_exhaustive]
 pub struct Local {}
 
+#[non_exhaustive]
+pub struct JwtClaims {
+    payload: String,
+    parsed: OnceCell<serde_json::Value>,
+    audience: OnceCell<Vec<String>>,
+}
+
+/// Authentication information for the caller of a reducer.
+pub struct AuthCtx {
+    is_internal: bool,
+    // I can't directly use a LazyLock without making this struct generic.
+    jwt: Box<dyn Deref<Target = Option<JwtClaims>>>,
+}
+
+impl AuthCtx {
+    fn new<F>(is_internal: bool, jwt_fn: F) -> Self
+    where
+        F: FnOnce() -> Option<JwtClaims> + 'static,
+    {
+        AuthCtx {
+            is_internal,
+            jwt: Box::new(LazyLock::new(jwt_fn)),
+        }
+    }
+
+    /// Create an AuthCtx for an internal call, with no JWT.
+    /// This represents a scheduled reducer.
+    pub fn internal() -> AuthCtx {
+        Self::new(true, || None)
+    }
+
+    /// Create an AuthCtx using the json claims from a JWT.
+    /// This can be used to write unit tests.
+    pub fn from_jwt_payload(jwt_payload: String) -> AuthCtx {
+        Self::new(false, move || Some(JwtClaims::new(jwt_payload)))
+    }
+
+    /// Create an AuthCtx that reads the JWT for the given connection id.
+    fn from_connection_id(connection_id: ConnectionId) -> AuthCtx {
+        Self::new(false, move || {
+            spacetimedb_bindings_sys::get_jwt(connection_id.as_le_byte_array()).map(JwtClaims::new)
+        })
+    }
+
+    /// True if this reducer was spawned from inside the database.
+    pub fn is_internal(&self) -> bool {
+        self.is_internal
+    }
+
+    /// Check if there is a JWT without loading it.
+    /// If is_internal is true, this will be false.
+    pub fn has_jwt(&self) -> bool {
+        self.jwt.is_some()
+    }
+
+    /// Load the jwt.
+    pub fn jwt(&self) -> Option<&JwtClaims> {
+        self.jwt.as_ref().deref().as_ref()
+    }
+}
+
+impl JwtClaims {
+    fn new(jwt: String) -> Self {
+        Self {
+            payload: jwt,
+            parsed: OnceCell::new(),
+            audience: OnceCell::new(),
+        }
+    }
+
+    fn get_parsed(&self) -> &serde_json::Value {
+        self.parsed
+            .get_or_init(|| serde_json::from_str(&self.payload).expect("Failed to parse JWT payload"))
+    }
+
+    pub fn subject(&self) -> &str {
+        // TODO: Add more error messages here.
+        self.get_parsed().get("sub").unwrap().as_str().unwrap()
+    }
+
+    pub fn issuer(&self) -> &str {
+        self.get_parsed().get("iss").unwrap().as_str().unwrap()
+    }
+
+    fn extract_audience(&self) -> Vec<String> {
+        let aud = self.get_parsed().get("aud").unwrap();
+        match aud {
+            serde_json::Value::String(s) => vec![s.clone()],
+            serde_json::Value::Array(arr) => arr.iter().filter_map(|v| v.as_str().map(String::from)).collect(),
+            _ => panic!("Unexpected type for 'aud' claim in JWT"),
+        }
+    }
+
+    pub fn audience(&self) -> &[String] {
+        self.audience.get_or_init(|| self.extract_audience())
+    }
+
+    // A convenience method, since this may not be in the token.
+    pub fn identity(&self) -> Identity {
+        Identity::from_claims(self.issuer(), self.subject())
+    }
+
+    // We can expose the whole payload for users that want to parse custom claims.
+    pub fn raw_payload(&self) -> &str {
+        &self.payload
+    }
+}
+
 // #[cfg(target_arch = "wasm32")]
 // #[global_allocator]
 // static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
@@ -900,3 +1033,37 @@ macro_rules! __volatile_nonatomic_schedule_immediate_impl {
         }
     };
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_single_audience() {
+        let example_payload = r#"
+        {
+          "iss": "https://securetoken.google.com/my-project-id",
+          "aud": "my-project-id",
+          "auth_time": 1695560000,
+          "user_id": "abc123XYZ",
+          "sub": "abc123XYZ",
+          "iat": 1695560100,
+          "exp": 1695563700,
+          "email": "user@example.com",
+          "email_verified": true,
+          "firebase": {
+            "identities": {
+              "email": ["user@example.com"]
+            },
+            "sign_in_provider": "password"
+          },
+          "name": "Jane Doe",
+          "picture": "https://lh3.googleusercontent.com/a-/profile.jpg"
+        }
+        "#;
+        let auth = AuthCtx::from_jwt_payload(example_payload.to_string());
+        let audience = auth.jwt().unwrap().audience();
+        assert_eq!(audience.len(), 1);
+        assert_eq!(audience, &["my-project-id".to_string()]);
+    }
+}

crates/bindings/src/rt.rs

@@ -488,13 +488,7 @@ extern "C" fn __call_reducer__(
 
     // Assemble the `ReducerContext`.
     let timestamp = Timestamp::from_micros_since_unix_epoch(timestamp as i64);
-    let ctx = ReducerContext {
-        db: crate::Local {},
-        sender,
-        timestamp,
-        connection_id: conn_id,
-        rng: std::cell::OnceCell::new(),
-    };
+    let ctx = ReducerContext::new(crate::Local {}, sender, conn_id, timestamp);
 
     // Fetch reducer function.
     let reducers = REDUCERS.get().unwrap();

crates/bindings/tests/snapshots/deps__spacetimedb_bindings_dependencies.snap

@@ -2,7 +2,7 @@
 source: crates/bindings/tests/deps.rs
 expression: "cargo tree -p spacetimedb -e no-dev --color never --target wasm32-unknown-unknown -f {lib}"
 ---
-total crates: 60
+total crates: 65
 spacetimedb
 ├── bytemuck
 ├── derive_more
@@ -21,6 +21,7 @@ spacetimedb
 ├── getrandom
 │   └── cfg_if
 ├── log
+├── once_cell
 ├── rand
 │   ├── rand_chacha
 │   │   ├── ppv_lite86
@@ -29,6 +30,11 @@ spacetimedb
 │   │       └── getrandom (*)
 │   └── rand_core (*)
 ├── scoped_tls
+├── serde_json
+│   ├── itoa
+│   ├── memchr
+│   ├── ryu
+│   └── serde
 ├── spacetimedb_bindings_macro
 │   ├── heck
 │   ├── humantime

crates/core/src/host/instance_env.rs

@@ -16,6 +16,7 @@ use spacetimedb_sats::{
 };
 use spacetimedb_table::indexes::RowPointer;
 use spacetimedb_table::table::RowRef;
+use std::fmt::Display;
 use std::ops::DerefMut;
 use std::sync::Arc;
 
@@ -499,6 +500,13 @@ impl From<GetTxError> for NodesError {
     }
 }
 
+impl Display for GetTxError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "not in a transaction")
+    }
+}
+impl std::error::Error for GetTxError {}
+
 #[cfg(test)]
 mod test {
     use super::*;

crates/core/src/host/mod.rs

@@ -145,6 +145,8 @@ pub enum AbiCall {
     ConsoleTimerStart,
     ConsoleTimerEnd,
     Identity,
+    JwtLength,
+    GetJwt,
 
     VolatileNonatomicScheduleImmediate,
 }

crates/core/src/host/wasm_common.rs

@@ -397,6 +397,7 @@ macro_rules! abi_funcs {
             "spacetime_10.0"::volatile_nonatomic_schedule_immediate,
 
             "spacetime_10.1"::bytes_source_remaining_length,
+            "spacetime_10.1"::get_jwt,
         }
     };
 }

crates/core/src/host/wasmtime/mod.rs

@@ -179,6 +179,18 @@ impl WasmPointee for spacetimedb_lib::Identity {
     }
 }
 
+impl WasmPointee for spacetimedb_lib::ConnectionId {
+    type Pointer = u32;
+    fn write_to(self, mem: &mut MemView, ptr: Self::Pointer) -> Result<(), MemError> {
+        let bytes = self.as_le_byte_array();
+        mem.deref_slice_mut(ptr, bytes.len() as u32)?.copy_from_slice(&bytes);
+        Ok(())
+    }
+    fn read_from(mem: &mut MemView, ptr: Self::Pointer) -> Result<Self, MemError> {
+        Ok(Self::from_le_byte_array(*mem.deref_array(ptr)?))
+    }
+}
+
 type WasmPtr<T> = <T as WasmPointee>::Pointer;
 
 /// Wraps access to WASM linear memory with some additional functionality.