storage: cleanup access/refresh token lookups

crates/data-model/src/lib.rs

@@ -44,7 +44,9 @@ pub use self::{
         AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, Client,
         InvalidRedirectUriError, JwksOrJwksUri, Pkce, Session, SessionState,
     },
-    tokens::{AccessToken, RefreshToken, TokenFormatError, TokenType},
+    tokens::{
+        AccessToken, AccessTokenState, RefreshToken, RefreshTokenState, TokenFormatError, TokenType,
+    },
     upstream_oauth2::{
         UpstreamOAuthAuthorizationSession, UpstreamOAuthAuthorizationSessionState,
         UpstreamOAuthLink, UpstreamOAuthProvider,

crates/data-model/src/tokens.rs

@@ -19,23 +19,133 @@ use rand::{distributions::Alphanumeric, Rng};
 use thiserror::Error;
 use ulid::Ulid;
 
+use crate::InvalidTransitionError;
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub enum AccessTokenState {
+    #[default]
+    Valid,
+    Revoked {
+        revoked_at: DateTime<Utc>,
+    },
+}
+
+impl AccessTokenState {
+    fn revoke(self, revoked_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        match self {
+            Self::Valid => Ok(Self::Revoked { revoked_at }),
+            Self::Revoked { .. } => Err(InvalidTransitionError),
+        }
+    }
+
+    /// Returns `true` if the refresh token state is [`Valid`].
+    ///
+    /// [`Valid`]: RefreshTokenState::Valid
+    #[must_use]
+    pub fn is_valid(&self) -> bool {
+        matches!(self, Self::Valid)
+    }
+
+    /// Returns `true` if the refresh token state is [`Revoked`].
+    ///
+    /// [`Revoked`]: RefreshTokenState::Revoked
+    #[must_use]
+    pub fn is_revoked(&self) -> bool {
+        matches!(self, Self::Revoked { .. })
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AccessToken {
     pub id: Ulid,
-    pub jti: String,
+    pub state: AccessTokenState,
+    pub session_id: Ulid,
     pub access_token: String,
     pub created_at: DateTime<Utc>,
     pub expires_at: DateTime<Utc>,
 }
 
+impl AccessToken {
+    #[must_use]
+    pub fn jti(&self) -> String {
+        self.id.to_string()
+    }
+
+    #[must_use]
+    pub fn is_valid(&self, now: DateTime<Utc>) -> bool {
+        self.state.is_valid() && self.expires_at > now
+    }
+
+    pub fn revoke(mut self, revoked_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.revoke(revoked_at)?;
+        Ok(self)
+    }
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub enum RefreshTokenState {
+    #[default]
+    Valid,
+    Consumed {
+        consumed_at: DateTime<Utc>,
+    },
+}
+
+impl RefreshTokenState {
+    fn consume(self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        match self {
+            Self::Valid => Ok(Self::Consumed { consumed_at }),
+            Self::Consumed { .. } => Err(InvalidTransitionError),
+        }
+    }
+
+    /// Returns `true` if the refresh token state is [`Valid`].
+    ///
+    /// [`Valid`]: RefreshTokenState::Valid
+    #[must_use]
+    pub fn is_valid(&self) -> bool {
+        matches!(self, Self::Valid)
+    }
+
+    /// Returns `true` if the refresh token state is [`Consumed`].
+    ///
+    /// [`Consumed`]: RefreshTokenState::Consumed
+    #[must_use]
+    pub fn is_consumed(&self) -> bool {
+        matches!(self, Self::Consumed { .. })
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RefreshToken {
     pub id: Ulid,
+    pub state: RefreshTokenState,
     pub refresh_token: String,
+    pub session_id: Ulid,
     pub created_at: DateTime<Utc>,
     pub access_token_id: Option<Ulid>,
 }
 
+impl std::ops::Deref for RefreshToken {
+    type Target = RefreshTokenState;
+
+    fn deref(&self) -> &Self::Target {
+        &self.state
+    }
+}
+
+impl RefreshToken {
+    #[must_use]
+    pub fn jti(&self) -> String {
+        self.id.to_string()
+    }
+
+    pub fn consume(mut self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.consume(consumed_at)?;
+        Ok(self)
+    }
+}
+
 /// Type of token to generate or validate
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum TokenType {