fix(cloud_auth): Ensure session is linked to user cork

dnys1 · dnys1 · commit 9b0cfc0d4962 · 2025-03-18T19:35:46.000Z
diff --git a/services/celest_cloud_auth/CHANGELOG.md b/services/celest_cloud_auth/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.1.0+2
 
 - fix: Ensure Cedar types are registered during seeding
+- fix: Ensure session is linked to user cork
 
 ## 0.1.0+1
 
diff --git a/services/celest_cloud_auth/lib/celest_cloud_auth.dart b/services/celest_cloud_auth/lib/celest_cloud_auth.dart
@@ -18,6 +18,7 @@ import 'package:celest_core/_internal.dart';
 import 'package:meta/meta.dart';
 import 'package:shelf/shelf.dart';
 
+export 'package:celest_cloud_auth/src/authentication/authentication_service.dart';
 export 'package:celest_cloud_auth/src/database/auth_database.dart';
 export 'package:celest_cloud_auth/src/email/email_provider.dart';
 export 'package:celest_cloud_auth/src/otp/otp_provider.dart';
diff --git a/services/celest_cloud_auth/lib/src/authentication/authentication_service.dart b/services/celest_cloud_auth/lib/src/authentication/authentication_service.dart
@@ -1,3 +1,5 @@
+import 'dart:typed_data';
+
 import 'package:cedar/cedar.dart';
 // ignore: invalid_use_of_internal_member
 import 'package:celest/src/runtime/http/cloud_middleware.dart';
@@ -253,13 +255,17 @@ extension type AuthenticationService._(_Deps _deps) implements Object {
   }
 
   Future<SessionStateSuccess> _authenticateUser({
+    required TypeId<Session> sessionId,
     required String userId,
   }) async {
     final user = await _db.getUser(userId: userId);
     if (user == null) {
       throw const NotFoundException('User not found');
     }
-    final cork = await _corks.createUserCork(user: user);
+    final cork = await _corks.createUserCork(
+      user: user,
+      sessionId: sessionId,
+    );
     return SessionStateSuccess(
       cork: cork,
       user: user,
@@ -268,6 +274,7 @@ extension type AuthenticationService._(_Deps _deps) implements Object {
   }
 
   Future<SessionStateSuccess> _createUser({
+    required TypeId<Session> sessionId,
     required AuthenticationFactor factor,
   }) async {
     final (user, cork) = await _db.transaction(() async {
@@ -289,7 +296,10 @@ extension type AuthenticationService._(_Deps _deps) implements Object {
           roles: const [EntityUid.of('Celest::Role', 'authenticated')],
         ),
       );
-      final cork = await _corks.createUserCork(user: user);
+      final cork = await _corks.createUserCork(
+        user: user,
+        sessionId: sessionId,
+      );
       return (user, cork);
     });
     return SessionStateSuccess(
@@ -350,10 +360,14 @@ extension type AuthenticationService._(_Deps _deps) implements Object {
         );
     }
     if (session.userId case final userId?) {
-      return _authenticateUser(userId: userId);
+      return _authenticateUser(
+        userId: userId,
+        sessionId: session.sessionId,
+      );
     }
     return _createUser(
       factor: factor,
+      sessionId: session.sessionId,
     );
   }
 
@@ -477,26 +491,35 @@ extension type AuthenticationService._(_Deps _deps) implements Object {
 
   @visibleForTesting
   Future<void> endSession({
-    required TypeId<Session> sessionId,
+    required String sessionId,
     required String sessionToken,
   }) async {
-    final session = await _db.authDrift
-        .getSession(sessionId: sessionId.encoded)
-        .getSingleOrNull();
-    if (session == null) {
-      throw const NotFoundException('Session not found');
+    if (sessionId.isNotEmpty) {
+      final session = await _db.authDrift
+          .getSession(sessionId: sessionId)
+          .getSingleOrNull();
+      if (session == null) {
+        throw const NotFoundException('Session not found');
+      }
+      sessionId = session.sessionId.encoded;
+    } else {
+      final sessionCork = CedarCork.parse(sessionToken);
+      await _corks.verify(cork: sessionCork);
+      sessionId = TypeId.decode(String.fromCharCodes(sessionCork.id)).encoded;
     }
-    final sessionCork = CedarCork.parse(sessionToken);
-    await _corks.verify(cork: sessionCork);
-
-    await _db.authDrift.deleteSession(sessionId: sessionId.encoded);
+    await _db.transaction(() async {
+      await _db.authDrift.deleteSession(sessionId: sessionId);
+      await _db.authDrift.deleteCork(
+        corkId: Uint8List.fromList(sessionId.codeUnits),
+      );
+    });
   }
 
   Future<Response> handleEndSession(Request request) async {
     final jsonRequest = await JsonUtf8.decodeStream(request.read());
     final pbRequest = pb.EndSessionRequest()..mergeFromProto3Json(jsonRequest);
     await endSession(
-      sessionId: TypeId.decode(pbRequest.sessionId),
+      sessionId: pbRequest.sessionId,
       sessionToken: pbRequest.sessionToken,
     );
     final response = pb.EndSessionResponse(
diff --git a/services/celest_cloud_auth/lib/src/authorization/corks_repository.dart b/services/celest_cloud_auth/lib/src/authorization/corks_repository.dart
@@ -42,33 +42,47 @@ extension type CorksRepository._(_Dependencies _deps) implements Object {
 
   /// Creates a new cork for the given [user].
   Future<Cork> createUserCork({
+    // TODO(dnys1): Make this required
+    TypeId<Session>? sessionId,
     required User user,
     EntityUid? audience,
   }) async {
-    final corkId = TypeId<Cork>().uuid.value;
+    sessionId ??= TypeId<Session>();
+    audience ??= issuer;
+    final cryptoKey = await _cryptoKeys.getOrMintHmacKey(
+      cryptoKeyId: sessionId.uuid.value,
+    );
+    final userId = EntityUid.of('Celest::User', user.userId);
+    final expireTime = DateTime.timestamp().add(const Duration(days: 30));
     final userEntity = Entity(
-      uid: EntityUid.of('Celest::User', user.userId),
+      uid: userId,
       parents: [
         ...user.roles,
         issuer,
       ],
-      attributes: RecordValue.fromJson(user.toJson()).attributes,
+      attributes: {
+        ...RecordValue.fromJson(user.toJson()).attributes,
+        'expireTime': Value.integer(expireTime.millisecondsSinceEpoch ~/ 1000),
+      },
     );
-    final corkBuilder = CedarCork.builder(corkId)
-      ..bearer = EntityUid.of('Celest::User', user.userId)
+    final corkBuilder = CedarCork.builder(sessionId.uuid.value)
+      ..bearer = userId
+      ..audience = audience
       ..issuer = issuer
-      ..audience = audience ?? issuer
       ..claims = userEntity;
-    final hmacKey = _cryptoKeys.rootKey;
-    final cork = await corkBuilder.build().sign(hmacKey.signer);
+    final cork = await corkBuilder.build().sign(cryptoKey.signer);
     await _db.transaction(() async {
       await _db.createEntity(userEntity);
-      await _db.authDrift.createCork(
-        corkId: corkId,
-        cryptoKeyId: hmacKey.cryptoKeyId,
+      await _db.authDrift.upsertCork(
+        corkId: cork.id,
+        cryptoKeyId: cryptoKey.cryptoKeyId,
         bearerType: 'Celest::User',
         bearerId: user.userId,
-        expireTime: DateTime.timestamp().add(const Duration(days: 30)),
+        audienceType: audience!.type,
+        audienceId: audience.id,
+        issuerType: issuer.type,
+        issuerId: issuer.id,
+        expireTime: expireTime,
       );
     });
     return cork;
@@ -97,7 +111,7 @@ extension type CorksRepository._(_Dependencies _deps) implements Object {
         ),
       },
     );
-    final corkBuilder = CedarCork.builder(cryptoKey.cryptoKeyId)
+    final corkBuilder = CedarCork.builder(session.sessionId.uuid.value)
       ..bearer = sessionUid
       ..issuer = issuer
       ..audience = issuer
@@ -110,6 +124,10 @@ extension type CorksRepository._(_Dependencies _deps) implements Object {
         cryptoKeyId: session.cryptoKeyId,
         bearerType: sessionUid.type,
         bearerId: sessionUid.id,
+        audienceType: issuer.type,
+        audienceId: issuer.id,
+        issuerType: issuer.type,
+        issuerId: issuer.id,
         expireTime: session.expireTime,
       );
     });
diff --git a/services/celest_cloud_auth/lib/src/database/schema/auth.drift b/services/celest_cloud_auth/lib/src/database/schema/auth.drift
@@ -315,6 +315,39 @@ createCork:
         :expire_time
     );
 
+upsertCork:
+    INSERT INTO corks(
+        cork_id, 
+        crypto_key_id, 
+        bearer_type, 
+        bearer_id, 
+        audience_type, 
+        audience_id, 
+        issuer_type, 
+        issuer_id, 
+        expire_time
+    )
+    VALUES (
+        :cork_id, 
+        :crypto_key_id, 
+        :bearer_type, 
+        :bearer_id, 
+        :audience_type, 
+        :audience_id, 
+        :issuer_type, 
+        :issuer_id, 
+        :expire_time
+    )
+    ON CONFLICT(cork_id) DO UPDATE SET
+        crypto_key_id = excluded.crypto_key_id,
+        bearer_type = excluded.bearer_type,
+        bearer_id = excluded.bearer_id,
+        audience_type = excluded.audience_type,
+        audience_id = excluded.audience_id,
+        issuer_type = excluded.issuer_type,
+        issuer_id = excluded.issuer_id,
+        expire_time = excluded.expire_time;
+
 -- Retrieves the cork for the given ID.
 getCork:
     SELECT * FROM corks
diff --git a/services/celest_cloud_auth/lib/src/database/schema/auth.drift.dart b/services/celest_cloud_auth/lib/src/database/schema/auth.drift.dart
@@ -2441,6 +2441,33 @@ class AuthDrift extends i7.ModularAccessor {
     );
   }
 
+  Future<int> upsertCork(
+      {required i2.Uint8List corkId,
+      required i2.Uint8List cryptoKeyId,
+      String? bearerType,
+      String? bearerId,
+      String? audienceType,
+      String? audienceId,
+      String? issuerType,
+      String? issuerId,
+      DateTime? expireTime}) {
+    return customInsert(
+      'INSERT INTO corks (cork_id, crypto_key_id, bearer_type, bearer_id, audience_type, audience_id, issuer_type, issuer_id, expire_time) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) ON CONFLICT (cork_id) DO UPDATE SET crypto_key_id = excluded.crypto_key_id, bearer_type = excluded.bearer_type, bearer_id = excluded.bearer_id, audience_type = excluded.audience_type, audience_id = excluded.audience_id, issuer_type = excluded.issuer_type, issuer_id = excluded.issuer_id, expire_time = excluded.expire_time',
+      variables: [
+        i0.Variable<i2.Uint8List>(corkId),
+        i0.Variable<i2.Uint8List>(cryptoKeyId),
+        i0.Variable<String>(bearerType),
+        i0.Variable<String>(bearerId),
+        i0.Variable<String>(audienceType),
+        i0.Variable<String>(audienceId),
+        i0.Variable<String>(issuerType),
+        i0.Variable<String>(issuerId),
+        i0.Variable<DateTime>(expireTime, const i5.TimestampType())
+      ],
+      updates: {corks},
+    );
+  }
+
   i0.Selectable<i3.Cork> getCork({required i2.Uint8List corkId}) {
     return customSelect('SELECT * FROM corks WHERE cork_id = ?1', variables: [
       i0.Variable<i2.Uint8List>(corkId)
diff --git a/services/celest_cloud_auth/test/authentication/authentication_service_test.dart b/services/celest_cloud_auth/test/authentication/authentication_service_test.dart
@@ -366,7 +366,7 @@ void main() {
         check(() => session.sessionToken).returnsNormally();
         await check(
           tester.authenticationService.endSession(
-            sessionId: session.sessionId,
+            sessionId: session.sessionId.encoded,
             sessionToken: session.sessionToken,
           ),
         ).completes();
@@ -402,7 +402,7 @@ void main() {
 
         await check(
           tester.authenticationService.endSession(
-            sessionId: session.sessionId,
+            sessionId: session.sessionId.encoded,
             sessionToken: cork.toString(),
           ),
         ).throws<InvalidSignatureException>();
diff --git a/services/celest_cloud_auth/test/tester.dart b/services/celest_cloud_auth/test/tester.dart
@@ -7,7 +7,6 @@ import 'package:celest/src/runtime/serve.dart';
 import 'package:celest_ast/celest_ast.dart';
 import 'package:celest_cloud/celest_cloud.dart' show CelestCloud, ClientType;
 import 'package:celest_cloud_auth/celest_cloud_auth.dart';
-import 'package:celest_cloud_auth/src/authentication/authentication_service.dart';
 import 'package:celest_cloud_auth/src/authorization/authorization_middleware.dart';
 import 'package:celest_cloud_auth/src/authorization/authorizer.dart';
 import 'package:celest_cloud_auth/src/authorization/cedar_interop.dart';
