Add AuthCtx to ReducerContext for rust

crates/auth/src/identity.rs

@@ -15,8 +15,7 @@ pub struct ConnectionAuthCtx {
 impl TryFrom<SpacetimeIdentityClaims> for ConnectionAuthCtx {
     type Error = anyhow::Error;
     fn try_from(claims: SpacetimeIdentityClaims) -> Result<Self, Self::Error> {
-        let payload =
-            serde_json::to_string(&claims).map_err(|e| anyhow::anyhow!("Failed to serialize claims: {}", e))?;
+        let payload = serde_json::to_string(&claims).map_err(|e| anyhow::anyhow!("Failed to serialize claims: {e}"))?;
         Ok(ConnectionAuthCtx {
             claims,
             jwt_payload: payload,
@@ -111,9 +110,7 @@ impl TryInto<SpacetimeIdentityClaims> for IncomingClaims {
         if let Some(token_identity) = self.identity {
             if token_identity != computed_identity {
                 return Err(anyhow::anyhow!(
-                    "Identity mismatch: token identity {:?} does not match computed identity {:?}",
-                    token_identity,
-                    computed_identity,
+                    "Identity mismatch: token identity {token_identity:?} does not match computed identity {computed_identity:?}",
                 ));
             }
         }

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,19 @@ 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
+        /// [`BytesSource`] to the given pointer.
+        /// If this is not found, [`BytesSource::INVALID`] (aka 0) will be written.
+        /// This must be called inside a transaction (because it reads from a system table).
+        ///
+        /// # Traps
+        ///
+        /// Traps if:
+        ///
+        /// - `connection_id` does not point to a valid little-endian `ConnectionId`.
+        /// - This is called outside a transaction.
+        pub fn get_jwt(connection_id_ptr: *const u8, bytes_source_id: *mut BytesSource);
     }
 
     /// What strategy does the database index use?
@@ -1118,6 +1132,19 @@ pub fn identity() -> [u8; 32] {
     buf
 }
 
+#[inline]
+pub fn get_jwt(connection_id: [u8; 16]) -> Option<raw::BytesSource> {
+    let mut source: raw::BytesSource = raw::BytesSource::INVALID;
+    unsafe {
+        raw::get_jwt(connection_id.as_ptr(), &mut source);
+    }
+    if source == raw::BytesSource::INVALID {
+        None // No JWT found.
+    } else {
+        Some(source)
+    }
+}
+
 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,112 @@ 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)
+    }
+
+    /// Creates 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)))
+    }
+
+    /// Creates an [`AuthCtx`] that reads the JWT for the given connection id.
+    fn from_connection_id(connection_id: ConnectionId) -> AuthCtx {
+        Self::new(false, move || rt::get_jwt(connection_id).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 +1031,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": "[email protected]",
+          "email_verified": true,
+          "firebase": {
+            "identities": {
+              "email": ["[email protected]"]
+            },
+            "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

@@ -2,6 +2,7 @@
 
 use crate::table::IndexAlgo;
 use crate::{sys, IterBuf, ReducerContext, ReducerResult, SpacetimeType, Table};
+use spacetimedb_bindings_sys::raw;
 pub use spacetimedb_lib::db::raw_def::v9::Lifecycle as LifecycleReducer;
 use spacetimedb_lib::db::raw_def::v9::{RawIndexAlgorithm, RawModuleDefV9Builder, TableType};
 use spacetimedb_lib::de::{self, Deserialize, Error as _, SeqProductAccess};
@@ -488,13 +489,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();
@@ -531,6 +526,20 @@ fn with_read_args<R>(args: BytesSource, logic: impl FnOnce(&[u8]) -> R) -> R {
 const NO_SPACE: u16 = errno::NO_SPACE.get();
 const NO_SUCH_BYTES: u16 = errno::NO_SUCH_BYTES.get();
 
+/// Look up the jwt associated with `connection_id`.
+pub fn get_jwt(connection_id: ConnectionId) -> Option<String> {
+    let mut buf = IterBuf::take();
+    let mut source: BytesSource = BytesSource::INVALID;
+    unsafe {
+        raw::get_jwt(connection_id.as_le_byte_array().as_ptr(), &mut source);
+    };
+    if source == BytesSource::INVALID {
+        return None;
+    }
+    read_bytes_source_into(source, &mut buf);
+    Some(std::str::from_utf8(&buf).unwrap().to_string())
+}
+
 /// Read `source` from the host fully into `buf`.
 fn read_bytes_source_into(source: BytesSource, buf: &mut Vec<u8>) {
     const INVALID: i16 = NO_SUCH_BYTES as i16;
@@ -565,8 +574,8 @@ fn read_bytes_source_into(source: BytesSource, buf: &mut Vec<u8>) {
         let buf_ptr = buf_ptr.as_mut_ptr().cast();
         let ret = unsafe { sys::raw::bytes_source_read(source, buf_ptr, &mut buf_len) };
         if ret <= 0 {
-            // SAFETY: `bytes_source_read` just appended `spare_len` bytes to `buf`.
-            unsafe { buf.set_len(buf.len() + spare_len) };
+            // SAFETY: `bytes_source_read` just appended `buf_len` bytes to `buf`.
+            unsafe { buf.set_len(buf.len() + buf_len) };
         }
         match ret {
             // Host side source exhausted, we're done.

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: 67
 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_core
 ├── spacetimedb_bindings_macro
 │   ├── heck
 │   ├── humantime
@@ -58,6 +64,7 @@ spacetimedb
 │   │   └── constant_time_eq
 │   │   [build-dependencies]
 │   │   └── cc
+│   │       ├── find_msvc_tools
 │   │       └── shlex
 │   ├── chrono
 │   │   └── num_traits
@@ -88,6 +95,7 @@ spacetimedb
 │   │   ├── enum_as_inner (*)
 │   │   ├── ethnum
 │   │   │   └── serde
+│   │   │       └── serde_core
 │   │   ├── hex
 │   │   ├── itertools (*)
 │   │   ├── second_stack