Add AuthCtx to ReducerContext for rust (#3288)

# Description of Changes

This exposes client credentials in reducer calls for rust.

# API and ABI breaking changes

API Changes:

The main API change is the addition of `AuthCtx` and the `sender_auth`
in `ReducerContext`. This also adds JwtClaims, which has some helpers
for getting commonly used claims.

ABI Changes:

This adds one new functions `get_jwt`. This uses
`st_connection_credentials` to look up the credentials associated with a
connection id.

This adds ABI version 10.2.

# Expected complexity level and risk

2. This adds new ABI functions

# Testing

I've done some manual testing with modified versions of the quickstart.
We should add some examples that use the new API.

Author: jsdt

Cargo.lock

@@ -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/auth/src/identity.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.
@@ -619,6 +620,31 @@ pub mod raw {
         pub fn bytes_source_remaining_length(source: BytesSource, out: *mut u32) -> i16;
     }
 
+    // See comment on previous `extern "C"` block re: ABI version.
+    #[link(wasm_import_module = "spacetime_10.2")]
+    extern "C" {
+        /// Finds the JWT payload associated with `connection_id`.
+        /// A `[ByteSourceId]` for the payload will be written to `target_ptr`.
+        /// If nothing is found for the connection, `[ByteSourceId::INVALID]` (zero) is written to `target_ptr`.
+        ///
+        /// This must be called inside a transaction (because it reads from a system table).
+        ///
+        /// # Errors
+        ///
+        /// Returns an error:
+        ///
+        /// - `NOT_IN_TRANSACTION`, when called outside of a transaction.
+        ///
+        /// # Traps
+        ///
+        /// Traps if:
+        ///
+        /// - `connection_id` does not point to a valid little-endian `ConnectionId`.
+        /// - `target_ptr` is NULL or `target_ptr[..size_of::<u32>()]` is not in bounds of WASM memory.
+        ///  - The `ByteSourceId` to be written to `target_ptr` would overflow [`u32::MAX`].
+        pub fn get_jwt(connection_id_ptr: *const u8, bytes_source_id: *mut BytesSource) -> u16;
+    }
+
     /// What strategy does the database index use?
     ///
     /// See also: <https://www.postgresql.org/docs/current/sql-createindex.html>
@@ -1118,6 +1144,29 @@ pub fn identity() -> [u8; 32] {
     buf
 }
 
+/// Finds the JWT payload associated with `connection_id`.
+/// If nothing is found for the connection, this returns None.
+/// If a payload is found, this will return a valid [`raw::BytesSource`].
+///
+/// This must be called inside a transaction (because it reads from a system table).
+///
+/// # Errors
+///
+/// This panics on any error. You can see details about errors in [`raw::get_jwt`].
+#[inline]
+pub fn get_jwt(connection_id: [u8; 16]) -> Option<raw::BytesSource> {
+    let source = unsafe {
+        call(|out| raw::get_jwt(connection_id.as_ptr(), out))
+            .unwrap_or_else(|errno: Errno| panic!("Error getting jwt: {errno}"))
+    };
+
+    if source == raw::BytesSource::INVALID {
+        None // No JWT found.
+    } else {
+        Some(source)
+    }
+}
+
 pub struct RowIter {
     raw: raw::RowIter,
 }

crates/bindings-sys/src/lib.rs

@@ -35,6 +35,7 @@ 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
 
 [dev-dependencies]
 insta.workspace = true

crates/bindings/Cargo.toml

@@ -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::LazyCell;
+use std::cell::{OnceCell, RefCell};
+use std::ops::Deref;
 
 #[cfg(feature = "unstable")]
 pub use client_visibility_filter::Filter;
@@ -751,6 +752,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>,
 }
@@ -763,10 +766,33 @@ 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(),
+        }
+    }
+
+    /// Returns the authorization information for the caller of this reducer.
+    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.
@@ -829,6 +855,115 @@ 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,
+    // NOTE(jsdt): cannot directly use a LazyLock without making this struct generic.
+    jwt: Box<dyn Deref<Target = Option<JwtClaims>>>,
+}
+
+impl AuthCtx {
+    fn new(is_internal: bool, jwt_fn: impl FnOnce() -> Option<JwtClaims> + 'static) -> Self {
+        AuthCtx {
+            is_internal,
+            jwt: Box::new(LazyCell::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))
+    }
+
+    /// Returns whether 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 [`AuthCtx::is_internal`] is true, this will return 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"))
+    }
+
+    /// Returns the tokens subject, from the sub claim.
+    pub fn subject(&self) -> &str {
+        self.get_parsed()
+            .get("sub")
+            .expect("Missing 'sub' claim")
+            .as_str()
+            .expect("Token 'sub' claim is not a string")
+    }
+
+    /// Returns the issuer for these credentials, from the iss claim.
+    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"),
+        }
+    }
+
+    /// Returns the audience for these credentials, from the aud claim.
+    pub fn audience(&self) -> &[String] {
+        self.audience.get_or_init(|| self.extract_audience())
+    }
+
+    /// Returns the identity for these credentials, which is
+    /// based on the iss and sub claims.
+    pub fn identity(&self) -> Identity {
+        Identity::from_claims(self.issuer(), self.subject())
+    }
+
+    /// Get the whole JWT payload as a json string.
+    pub fn raw_payload(&self) -> &str {
+        &self.payload
+    }
+}
 /// The read-only version of [`Local`]
 #[non_exhaustive]
 pub struct LocalReadOnly {}
@@ -937,3 +1072,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/lib.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();
@@ -531,6 +525,17 @@ 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 source = sys::get_jwt(connection_id.as_le_byte_array())?;
+    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 +570,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/src/rt.rs

@@ -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: 66
 spacetimedb
 ├── bytemuck
 ├── derive_more
@@ -29,6 +29,11 @@ spacetimedb
 │   │       └── getrandom (*)
 │   └── rand_core (*)
 ├── scoped_tls
+├── serde_json
+│   ├── itoa
+│   ├── memchr
+│   ├── ryu
+│   └── serde_core
 ├── spacetimedb_bindings_macro
 │   ├── heck
 │   ├── humantime
@@ -64,6 +69,7 @@ spacetimedb
 │   │   └── constant_time_eq
 │   │   [build-dependencies]
 │   │   └── cc
+│   │       ├── find_msvc_tools
 │   │       └── shlex
 │   ├── chrono
 │   │   └── num_traits
@@ -90,6 +96,7 @@ spacetimedb
 │   │   ├── enum_as_inner (*)
 │   │   ├── ethnum
 │   │   │   └── serde
+│   │   │       └── serde_core
 │   │   ├── hex
 │   │   ├── itertools (*)
 │   │   ├── second_stack