Add abi for credentials

Author: jsdt

crates/bindings-sys/src/lib.rs

@@ -588,6 +588,15 @@ pub mod raw {
         ///
         /// - `out_ptr` is NULL or `out` is not in bounds of WASM memory.
         pub fn identity(out_ptr: *mut u8);
+
+        /// Check the size of the jwt associated with the given connection.
+        /// Returns 0 if there is no jwt for the connection.
+        pub fn jwt_len(connection_id_ptr: *const u8, out_ptr: *mut u32);
+
+        /// Write the jwt payload for the given connection id to the out_ptr.
+        /// target_ptr_len will be set to the number of bytes written to the buffer.
+        /// If the buffer is too small, or the connection has no jwt, target_ptr_len will be set to 0.
+        pub fn get_jwt(connection_id_ptr: *const u8, target_ptr: *mut u8, target_ptr_len: *mut u32);
     }
 
     /// What strategy does the database index use?
@@ -1089,6 +1098,34 @@ pub fn identity() -> [u8; 32] {
     buf
 }
 
+#[inline]
+pub fn jwt_length(connection_id: [u8; 16]) -> Option<u32> {
+    let mut v: u32 = 0;
+    unsafe { raw::jwt_len(connection_id.as_ptr(), &mut v) }
+    if v == 0 {
+        None
+    } else {
+        Some(v)
+    }
+}
+
+#[inline]
+pub fn get_jwt(connection_id: [u8; 16]) -> Option<String> {
+    let Some(jwt_len) = jwt_length(connection_id) else {
+        return None; // No JWT found.
+    };
+    let mut buf = vec![0u8; jwt_len as usize];
+
+    let mut v: u32 = buf.len() as u32;
+    unsafe {
+        raw::get_jwt(connection_id.as_ptr(), buf.as_mut_ptr(), &mut v);
+    }
+    if v == 0 {
+        return None; // No JWT found.
+    }
+    Some(std::str::from_utf8(&buf[..v as usize]).unwrap().to_string())
+}
+
 pub struct RowIter {
     raw: raw::RowIter,
 }

crates/core/src/host/mod.rs

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

crates/core/src/host/wasm_common.rs

@@ -392,6 +392,8 @@ macro_rules! abi_funcs {
             "spacetime_10.0"::datastore_btree_scan_bsatn,
             "spacetime_10.0"::datastore_delete_by_btree_scan_bsatn,
             "spacetime_10.0"::identity,
+            "spacetime_10.0"::get_jwt,
+            "spacetime_10.0"::jwt_len,
 
             // unstable:
             "spacetime_10.0"::volatile_nonatomic_schedule_immediate,

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.

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

@@ -12,7 +12,7 @@ use crate::host::wasm_common::{
 };
 use crate::host::AbiCall;
 use anyhow::Context as _;
-use spacetimedb_lib::Timestamp;
+use spacetimedb_lib::{ConnectionId, Timestamp};
 use spacetimedb_primitives::{errno, ColId};
 use wasmtime::{AsContext, Caller, StoreContextMut};
 
@@ -22,6 +22,7 @@ use super::{Mem, MemView, NullableMemOp, WasmError, WasmPointee, WasmPtr};
 use instrumentation::noop as span;
 #[cfg(feature = "spacetimedb-wasm-instance-env-times")]
 use instrumentation::op as span;
+use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
 
 /// A `WasmInstanceEnv` provides the connection between a module
 /// and the database.
@@ -1208,6 +1209,69 @@ impl WasmInstanceEnv {
         })
     }
 
+    pub fn jwt_len(
+        caller: Caller<'_, Self>,
+        connection_id: WasmPtr<ConnectionId>,
+        out_ptr: WasmPtr<u32>,
+    ) -> RtResult<()> {
+        log::info!("Calling has_jwt");
+        Self::with_span(caller, AbiCall::JwtLength, |caller| {
+            let (mem, env) = Self::mem_env(caller);
+            let cid = ConnectionId::read_from(mem, connection_id)?;
+            let length = env
+                .instance_env
+                .tx
+                .get()
+                .unwrap()
+                .get_jwt_payload(cid)?
+                .map(|p| p.len() as u32)
+                .unwrap_or(0u32);
+            length.write_to(mem, out_ptr)?;
+            Ok(())
+        })
+    }
+
+    pub fn get_jwt(
+        caller: Caller<'_, Self>,
+        connection_id: WasmPtr<ConnectionId>,
+        target_ptr: WasmPtr<u8>,
+        target_ptr_len: WasmPtr<u32>,
+    ) -> RtResult<()> {
+        log::info!("Calling get_jwt");
+        Self::with_span(caller, AbiCall::GetJwt, |caller| {
+            let (mem, env) = Self::mem_env(caller);
+            let cid = ConnectionId::read_from(mem, connection_id)?;
+            let jwt = match env.instance_env.tx.get().unwrap().get_jwt_payload(cid)? {
+                None => {
+                    // Consider logging here, since this should only happen during an upgrade.
+                    0u32.write_to(mem, target_ptr_len)?;
+                    return Ok(());
+                }
+                Some(jwt) => jwt,
+            };
+            log::info!("JWT payload found for connection ID: {:?}: {}", cid.to_hex(), jwt);
+            let jwt_len = jwt.len();
+            // Read `buffer_len`, i.e., the capacity of `buffer` pointed to by `buffer_ptr`.
+            let buffer_len = u32::read_from(mem, target_ptr_len)?;
+            log::info!("buffer_len: {buffer_len}");
+            if buffer_len < jwt_len as u32 {
+                return Err(anyhow::anyhow!("buffer too small to hold JWT payload"));
+            }
+            log::info!("About to write length");
+            // Write the length of the JWT payload to the target pointer.
+            (jwt_len as u32).write_to(mem, target_ptr_len)?;
+            log::info!("wrote length of {jwt_len}");
+            log::info!("Byte len {}", jwt.len());
+
+            // Write the JWT payload to the target pointer.
+            // Get a mutable view to the `buffer`.
+            let buffer = mem.deref_slice_mut(target_ptr, buffer_len)?;
+            buffer[..jwt_len].copy_from_slice(jwt.as_bytes());
+            log::info!("wrote jwt bytes to slice");
+            Ok(())
+        })
+    }
+
     /// Writes the identity of the module into `out = out_ptr[..32]`.
     ///
     /// # Traps