feat(user-management): add session (#185)

Author: DanielleHuisman

Cargo.lock

@@ -13,7 +13,10 @@ native-tls = ["reqwest/native-tls"]
 rustls-tls = ["reqwest/rustls-tls"]
 
 [dependencies]
+aead = { version = "0.5.2", features = ["std"] }
+aes-gcm = "0.10.3"
 async-trait = "0.1.88"
+base64 = "0.22.1"
 chrono = { version = "0.4.40", features = ["serde"] }
 derive_more = { version = "2.0.1", features = ["deref", "display", "from"] }
 jsonwebtoken = { version = "10.0.0", features = ["rust_crypto"] }
@@ -28,6 +31,7 @@ urlencoding = "2.1.3"
 [dev-dependencies]
 matches = "0.1.10"
 mockito = "1.0.0"
+rsa = "0.9.8"
 tokio = { version = "1.44.2", default-features = false, features = [
     "macros",
     "rt-multi-thread",

Cargo.toml

@@ -2,6 +2,7 @@ mod api_key;
 mod metadata;
 mod paginated_list;
 mod pagination_params;
+mod remote_jwk_set;
 mod timestamps;
 mod unpaginated_list;
 mod url_encodable_vec;
@@ -10,6 +11,7 @@ pub use api_key::*;
 pub use metadata::*;
 pub use paginated_list::*;
 pub use pagination_params::*;
+pub use remote_jwk_set::*;
 pub use timestamps::*;
 pub use unpaginated_list::*;
 pub(crate) use url_encodable_vec::*;

src/core/types.rs

@@ -0,0 +1,84 @@
+use std::sync::{Arc, Mutex};
+
+use chrono::{DateTime, FixedOffset, TimeDelta, Utc};
+use jsonwebtoken::jwk::{Jwk, JwkSet};
+use thiserror::Error;
+use url::Url;
+
+use crate::{ResponseExt, WorkOsError, WorkOsResult, user_management::GetJwksError};
+
+type Entry = (JwkSet, DateTime<FixedOffset>);
+
+/// An error returned from [`RemoteJwkSet::find`].
+#[derive(Debug, Error)]
+pub enum FindJwkError {
+    /// Get JWKS error.
+    #[error(transparent)]
+    GetJwks(WorkOsError<GetJwksError>),
+
+    /// Poison error.
+    #[error("poison error: {0}")]
+    Poison(String),
+}
+
+impl From<FindJwkError> for WorkOsError<FindJwkError> {
+    fn from(value: FindJwkError) -> Self {
+        Self::Operation(value)
+    }
+}
+
+/// Remote JSON Web Key Set (JWKS).
+#[derive(Clone)]
+pub struct RemoteJwkSet {
+    client: reqwest::Client,
+    url: Url,
+    jwks: Arc<Mutex<Option<Entry>>>,
+}
+
+impl RemoteJwkSet {
+    pub(crate) fn new(client: reqwest::Client, url: Url) -> Self {
+        RemoteJwkSet {
+            client,
+            url,
+            jwks: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    /// Find the key in the set that matches the given key id, if any.
+    pub async fn find(&self, kid: &str) -> WorkOsResult<Option<Jwk>, FindJwkError> {
+        {
+            let jwks = self
+                .jwks
+                .lock()
+                .map_err(|err| FindJwkError::Poison(err.to_string()))?;
+
+            if let Some((jwks, expires_at)) = jwks.as_ref()
+                && *expires_at > Utc::now().fixed_offset()
+            {
+                return Ok(jwks.find(kid).cloned());
+            }
+        }
+
+        let new_jwks = self
+            .client
+            .get(self.url.as_str())
+            .send()
+            .await?
+            .handle_unauthorized_or_generic_error()
+            .await?
+            .json::<JwkSet>()
+            .await?;
+
+        let key = new_jwks.find(kid).cloned();
+
+        // TODO: Consider using the expiry of keys instead?
+        let mut jwks = self
+            .jwks
+            .lock()
+            .map_err(|err| FindJwkError::Poison(err.to_string()))?;
+
+        *jwks = Some((new_jwks, Utc::now().fixed_offset() + TimeDelta::minutes(5)));
+
+        Ok(key)
+    }
+}

src/core/types/remote_jwk_set.rs

@@ -105,10 +105,12 @@ impl GetAuthorizationUrl for Sso<'_> {
                 ),
             };
 
+            let redirect_uri = urlencoding::encode(redirect_uri);
+
             let mut query_params: querystring::QueryParams = vec![
                 ("response_type", "code"),
                 ("client_id", &client_id),
-                ("redirect_uri", redirect_uri),
+                ("redirect_uri", &redirect_uri),
                 (connection_selector_param.0, &connection_selector_param.1),
             ];
 
@@ -149,7 +151,7 @@ mod test {
         assert_eq!(
             authorization_url,
             Url::parse(
-                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https://your-app.com/callback&connection=conn_1234"
+                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https%3A%2F%2Fyour-app.com%2Fcallback&connection=conn_1234"
             )
             .unwrap()
         )
@@ -174,7 +176,7 @@ mod test {
         assert_eq!(
             authorization_url,
             Url::parse(
-                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https://your-app.com/callback&organization=org_1234"
+                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https%3A%2F%2Fyour-app.com%2Fcallback&organization=org_1234"
             )
             .unwrap()
         )
@@ -197,7 +199,7 @@ mod test {
         assert_eq!(
             authorization_url,
             Url::parse(
-                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https://your-app.com/callback&provider=GoogleOAuth"
+                "https://api.workos.com/sso/authorize?response_type=code&client_id=client_123456789&redirect_uri=https%3A%2F%2Fyour-app.com%2Fcallback&provider=GoogleOAuth"
             )
             .unwrap()
         )

src/sso/operations/get_authorization_url.rs

@@ -2,24 +2,81 @@
 //!
 //! [WorkOS Docs: User Management](https://workos.com/docs/user-management)
 
+mod cookie_session;
 mod operations;
 mod types;
 
+use std::sync::{Arc, Mutex};
+
+pub use cookie_session::*;
 pub use operations::*;
+use thiserror::Error;
 pub use types::*;
 
-use crate::WorkOs;
+use crate::{RemoteJwkSet, WorkOs};
+
+/// An error returned from [`UserManagement::jwks`].
+#[derive(Debug, Error)]
+pub enum JwksError {
+    /// Missing client ID
+    #[error("missing client ID")]
+    MissingClientId,
+
+    /// Poison error.
+    #[error("poison error: {0}")]
+    Poison(String),
+
+    /// URL error.
+    #[error(transparent)]
+    Url(#[from] url::ParseError),
+}
 
 /// User Management.
 ///
 /// [WorkOS Docs: User Management](https://workos.com/docs/user-management)
 pub struct UserManagement<'a> {
     workos: &'a WorkOs,
+    jwks: Arc<Mutex<Option<RemoteJwkSet>>>,
 }
 
 impl<'a> UserManagement<'a> {
     /// Returns a new [`UserManagement`] instance for the provided WorkOS client.
     pub fn new(workos: &'a WorkOs) -> Self {
-        Self { workos }
+        Self {
+            workos,
+            jwks: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    /// Get remote JSON Web Key Set (JWKS).
+    pub fn jwks(&'a self) -> Result<RemoteJwkSet, JwksError> {
+        let mut jwks = self
+            .jwks
+            .lock()
+            .map_err(|err| JwksError::Poison(err.to_string()))?;
+
+        if let Some(jwks) = jwks.as_ref() {
+            return Ok(jwks.clone());
+        }
+
+        let Some(client_id) = self.workos.client_id() else {
+            return Err(JwksError::MissingClientId);
+        };
+
+        let new_jwks =
+            RemoteJwkSet::new(self.workos.client().clone(), self.get_jwks_url(client_id)?);
+
+        *jwks = Some(new_jwks.clone());
+
+        Ok(new_jwks)
+    }
+
+    /// Load the session by providing the sealed session and the cookie password.
+    pub fn load_sealed_session(
+        &'a self,
+        session_data: &'a str,
+        cookie_password: &'a str,
+    ) -> CookieSession<'a> {
+        CookieSession::new(self, session_data, cookie_password)
     }
 }