Add id column on tokens and sessions

nexus/auth/src/authn/external/session_cookie.rs

@@ -13,13 +13,15 @@ use chrono::{DateTime, Duration, Utc};
 use dropshot::HttpError;
 use http::HeaderValue;
 use nexus_types::authn::cookies::parse_cookies;
+use omicron_uuid_kinds::{ConsoleSessionKind, TypedUuid};
 use slog::debug;
 use uuid::Uuid;
 
 // many parts of the implementation will reference this OWASP guide
 // https://cheatsheetseries.owasp.org/cheatsheets/Session_Management_Cheat_Sheet.html
 
 pub trait Session {
+    fn id(&self) -> TypedUuid<ConsoleSessionKind>;
     fn silo_user_id(&self) -> Uuid;
     fn silo_id(&self) -> Uuid;
     fn time_last_used(&self) -> DateTime<Utc>;
@@ -39,11 +41,14 @@ pub trait SessionStore {
     /// Extend session by updating time_last_used to now
     async fn session_update_last_used(
         &self,
-        token: String,
+        id: TypedUuid<ConsoleSessionKind>,
     ) -> Option<Self::SessionModel>;
 
     /// Mark session expired
-    async fn session_expire(&self, token: String) -> Option<()>;
+    async fn session_expire(
+        &self,
+        id: TypedUuid<ConsoleSessionKind>,
+    ) -> Option<()>;
 
     /// Maximum time session can remain idle before expiring
     fn session_idle_timeout(&self) -> Duration;
@@ -131,7 +136,7 @@ where
         // expired
         let now = Utc::now();
         if session.time_last_used() + ctx.session_idle_timeout() < now {
-            let expired_session = ctx.session_expire(token.clone()).await;
+            let expired_session = ctx.session_expire(session.id()).await;
             if expired_session.is_none() {
                 debug!(log, "failed to expire session")
             }
@@ -151,7 +156,7 @@ where
         // existed longer than absolute_timeout, it is expired and we can no
         // longer extend the session
         if session.time_created() + ctx.session_absolute_timeout() < now {
-            let expired_session = ctx.session_expire(token.clone()).await;
+            let expired_session = ctx.session_expire(session.id()).await;
             if expired_session.is_none() {
                 debug!(log, "failed to expire session")
             }
@@ -172,7 +177,7 @@ where
         // authenticated for this request at this point. The next request might
         // be wrongly considered idle, but that's a problem for the next
         // request.
-        let updated_session = ctx.session_update_last_used(token).await;
+        let updated_session = ctx.session_update_last_used(session.id()).await;
         if updated_session.is_none() {
             debug!(log, "failed to extend session")
         }
@@ -199,26 +204,32 @@ mod test {
     use async_trait::async_trait;
     use chrono::{DateTime, Duration, Utc};
     use http;
+    use omicron_uuid_kinds::ConsoleSessionKind;
+    use omicron_uuid_kinds::TypedUuid;
     use slog;
-    use std::collections::HashMap;
     use std::sync::Mutex;
     use uuid::Uuid;
 
     // the mutex is annoying, but we need it in order to mutate the hashmap
     // without passing TestServerContext around as mutable
     struct TestServerContext {
-        sessions: Mutex<HashMap<String, FakeSession>>,
+        sessions: Mutex<Vec<FakeSession>>,
     }
 
-    #[derive(Clone, Copy)]
+    #[derive(Clone)]
     struct FakeSession {
+        id: TypedUuid<ConsoleSessionKind>,
+        token: String,
         silo_user_id: Uuid,
         silo_id: Uuid,
         time_created: DateTime<Utc>,
         time_last_used: DateTime<Utc>,
     }
 
     impl Session for FakeSession {
+        fn id(&self) -> TypedUuid<ConsoleSessionKind> {
+            self.id
+        }
         fn silo_user_id(&self) -> Uuid {
             self.silo_user_id
         }
@@ -241,23 +252,37 @@ mod test {
             &self,
             token: String,
         ) -> Option<Self::SessionModel> {
-            self.sessions.lock().unwrap().get(&token).map(|s| *s)
+            self.sessions
+                .lock()
+                .unwrap()
+                .iter()
+                .find(|s| s.token == token)
+                .map(|s| s.clone())
         }
 
         async fn session_update_last_used(
             &self,
-            token: String,
+            id: TypedUuid<ConsoleSessionKind>,
         ) -> Option<Self::SessionModel> {
             let mut sessions = self.sessions.lock().unwrap();
-            let session = *sessions.get(&token).unwrap();
-            let new_session =
-                FakeSession { time_last_used: Utc::now(), ..session };
-            (*sessions).insert(token, new_session)
+            if let Some(pos) = sessions.iter().position(|s| s.id == id) {
+                let new_session = FakeSession {
+                    time_last_used: Utc::now(),
+                    ..sessions[pos].clone()
+                };
+                sessions[pos] = new_session.clone();
+                Some(new_session)
+            } else {
+                None
+            }
         }
 
-        async fn session_expire(&self, token: String) -> Option<()> {
+        async fn session_expire(
+            &self,
+            id: TypedUuid<ConsoleSessionKind>,
+        ) -> Option<()> {
             let mut sessions = self.sessions.lock().unwrap();
-            (*sessions).remove(&token);
+            sessions.retain(|s| s.id != id);
             Some(())
         }
 
@@ -295,32 +320,29 @@ mod test {
 
     #[tokio::test]
     async fn test_missing_cookie() {
-        let context =
-            TestServerContext { sessions: Mutex::new(HashMap::new()) };
+        let context = TestServerContext { sessions: Mutex::new(Vec::new()) };
         let result = authn_with_cookie(&context, None).await;
         assert!(matches!(result, SchemeResult::NotRequested));
     }
 
     #[tokio::test]
     async fn test_other_cookie() {
-        let context =
-            TestServerContext { sessions: Mutex::new(HashMap::new()) };
+        let context = TestServerContext { sessions: Mutex::new(Vec::new()) };
         let result = authn_with_cookie(&context, Some("other=def")).await;
         assert!(matches!(result, SchemeResult::NotRequested));
     }
 
     #[tokio::test]
     async fn test_expired_cookie_idle() {
         let context = TestServerContext {
-            sessions: Mutex::new(HashMap::from([(
-                "abc".to_string(),
-                FakeSession {
-                    silo_user_id: Uuid::new_v4(),
-                    silo_id: Uuid::new_v4(),
-                    time_last_used: Utc::now() - Duration::hours(2),
-                    time_created: Utc::now() - Duration::hours(2),
-                },
-            )])),
+            sessions: Mutex::new(vec![FakeSession {
+                id: TypedUuid::new_v4(),
+                token: "abc".to_string(),
+                silo_user_id: Uuid::new_v4(),
+                silo_id: Uuid::new_v4(),
+                time_last_used: Utc::now() - Duration::hours(2),
+                time_created: Utc::now() - Duration::hours(2),
+            }]),
         };
         let result = authn_with_cookie(&context, Some("session=abc")).await;
         assert!(matches!(
@@ -332,21 +354,21 @@ mod test {
         ));
 
         // key should be removed from sessions dict, i.e., session deleted
-        assert!(!context.sessions.lock().unwrap().contains_key("abc"))
+        let sessions = context.sessions.lock().unwrap();
+        assert!(!sessions.iter().any(|s| s.token == "abc"))
     }
 
     #[tokio::test]
     async fn test_expired_cookie_absolute() {
         let context = TestServerContext {
-            sessions: Mutex::new(HashMap::from([(
-                "abc".to_string(),
-                FakeSession {
-                    silo_user_id: Uuid::new_v4(),
-                    silo_id: Uuid::new_v4(),
-                    time_last_used: Utc::now(),
-                    time_created: Utc::now() - Duration::hours(20),
-                },
-            )])),
+            sessions: Mutex::new(vec![FakeSession {
+                id: TypedUuid::new_v4(),
+                token: "abc".to_string(),
+                silo_user_id: Uuid::new_v4(),
+                silo_id: Uuid::new_v4(),
+                time_last_used: Utc::now(),
+                time_created: Utc::now() - Duration::hours(20),
+            }]),
         };
         let result = authn_with_cookie(&context, Some("session=abc")).await;
         assert!(matches!(
@@ -359,22 +381,21 @@ mod test {
 
         // key should be removed from sessions dict, i.e., session deleted
         let sessions = context.sessions.lock().unwrap();
-        assert!(!sessions.contains_key("abc"))
+        assert!(!sessions.iter().any(|s| s.token == "abc"))
     }
 
     #[tokio::test]
     async fn test_valid_cookie() {
         let time_last_used = Utc::now() - Duration::seconds(5);
         let context = TestServerContext {
-            sessions: Mutex::new(HashMap::from([(
-                "abc".to_string(),
-                FakeSession {
-                    silo_user_id: Uuid::new_v4(),
-                    silo_id: Uuid::new_v4(),
-                    time_last_used,
-                    time_created: Utc::now(),
-                },
-            )])),
+            sessions: Mutex::new(vec![FakeSession {
+                id: TypedUuid::new_v4(),
+                token: "abc".to_string(),
+                silo_user_id: Uuid::new_v4(),
+                silo_id: Uuid::new_v4(),
+                time_last_used,
+                time_created: Utc::now(),
+            }]),
         };
         let result = authn_with_cookie(&context, Some("session=abc")).await;
         assert!(matches!(
@@ -384,13 +405,13 @@ mod test {
 
         // valid cookie should have updated time_last_used
         let sessions = context.sessions.lock().unwrap();
-        assert!(sessions.get("abc").unwrap().time_last_used > time_last_used)
+        let session = sessions.iter().find(|s| s.token == "abc").unwrap();
+        assert!(session.time_last_used > time_last_used)
     }
 
     #[tokio::test]
     async fn test_garbage_cookie() {
-        let context =
-            TestServerContext { sessions: Mutex::new(HashMap::new()) };
+        let context = TestServerContext { sessions: Mutex::new(Vec::new()) };
         let result =
             authn_with_cookie(&context, Some("unparsable garbage!!!!!1")).await;
         assert!(matches!(result, SchemeResult::NotRequested));

nexus/auth/src/authz/api_resources.rs

@@ -941,7 +941,7 @@ authz_resource! {
 authz_resource! {
     name = "ConsoleSession",
     parent = "Fleet",
-    primary_key = String,
+    primary_key = { uuid_kind = ConsoleSessionKind },
     roles_allowed = false,
     polar_snippet = FleetChild,
 }
@@ -957,7 +957,7 @@ authz_resource! {
 authz_resource! {
     name = "DeviceAccessToken",
     parent = "Fleet",
-    primary_key = String, // token
+    primary_key = { uuid_kind = AccessTokenKind },
     roles_allowed = false,
     polar_snippet = FleetChild,
 }

nexus/auth/src/context.rs

@@ -11,6 +11,8 @@ use crate::authz::AuthorizedResource;
 use crate::storage::Storage;
 use chrono::{DateTime, Utc};
 use omicron_common::api::external::Error;
+use omicron_uuid_kinds::ConsoleSessionKind;
+use omicron_uuid_kinds::TypedUuid;
 use slog::debug;
 use slog::o;
 use slog::trace;
@@ -352,6 +354,10 @@ impl OpContext {
 }
 
 impl Session for ConsoleSessionWithSiloId {
+    fn id(&self) -> TypedUuid<ConsoleSessionKind> {
+        self.console_session.id()
+    }
+
     fn silo_user_id(&self) -> Uuid {
         self.console_session.silo_user_id
     }

nexus/db-lookup/src/lookup.rs

@@ -26,6 +26,8 @@ use nexus_types::identity::Resource;
 use omicron_common::api::external::Error;
 use omicron_common::api::external::InternalContext;
 use omicron_common::api::external::{LookupResult, LookupType, ResourceType};
+use omicron_uuid_kinds::AccessTokenKind;
+use omicron_uuid_kinds::ConsoleSessionKind;
 use omicron_uuid_kinds::PhysicalDiskUuid;
 use omicron_uuid_kinds::SupportBundleUuid;
 use omicron_uuid_kinds::TufArtifactKind;
@@ -199,19 +201,12 @@ impl<'a> LookupPath<'a> {
 
     // Fleet-level resources
 
-    /// Select a resource of type ConsoleSession, identified by its `token`
-    pub fn console_session_token<'b, 'c>(
+    /// Select a resource of type ConsoleSession, identified by its `id`
+    pub fn console_session_id(
         self,
-        token: &'b str,
-    ) -> ConsoleSession<'c>
-    where
-        'a: 'c,
-        'b: 'c,
-    {
-        ConsoleSession::PrimaryKey(
-            Root { lookup_root: self },
-            token.to_string(),
-        )
+        id: TypedUuid<ConsoleSessionKind>,
+    ) -> ConsoleSession<'a> {
+        ConsoleSession::PrimaryKey(Root { lookup_root: self }, id)
     }
 
     /// Select a resource of type DeviceAuthRequest, identified by its `user_code`
@@ -229,19 +224,12 @@ impl<'a> LookupPath<'a> {
         )
     }
 
-    /// Select a resource of type DeviceAccessToken, identified by its `token`
-    pub fn device_access_token<'b, 'c>(
+    /// Select a resource of type DeviceAccessToken, identified by its `id`
+    pub fn device_access_token_id(
         self,
-        token: &'b str,
-    ) -> DeviceAccessToken<'c>
-    where
-        'a: 'c,
-        'b: 'c,
-    {
-        DeviceAccessToken::PrimaryKey(
-            Root { lookup_root: self },
-            token.to_string(),
-        )
+        id: TypedUuid<AccessTokenKind>,
+    ) -> DeviceAccessToken<'a> {
+        DeviceAccessToken::PrimaryKey(Root { lookup_root: self }, id)
     }
 
     /// Select a resource of type RoleBuiltin, identified by its `name`
@@ -761,9 +749,7 @@ lookup_resource! {
     ancestors = [],
     lookup_by_name = false,
     soft_deletes = false,
-    primary_key_columns = [
-        { column_name = "token", rust_type = String },
-    ]
+    primary_key_columns = [ { column_name = "id", uuid_kind = ConsoleSessionKind } ]
 }
 
 lookup_resource! {
@@ -781,9 +767,7 @@ lookup_resource! {
     ancestors = [],
     lookup_by_name = false,
     soft_deletes = false,
-    primary_key_columns = [
-        { column_name = "token", rust_type = String },
-    ]
+    primary_key_columns = [ { column_name = "id", uuid_kind = AccessTokenKind } ]
 }
 
 lookup_resource! {