From 1f5fa5b1db0e66a100d669e2bdd0e58a8ad9298c Mon Sep 17 00:00:00 2001
From: Paul Kraft <pauljohanneskraft@users.noreply.github.com>
Date: Thu, 7 Nov 2024 12:50:42 -0800
Subject: [PATCH] Return customClaims in both beforeUserCreated and
 beforeUserSignedIn

---
 .../models/src/types/userRegistration.ts      | 19 ++++++++++++++++++-
 functions/src/functions/blocking.ts           | 16 +++++++++++-----
 .../services/user/databaseUserService.test.ts | 12 +++---------
 .../src/services/user/databaseUserService.ts  | 18 ------------------
 .../src/services/user/userService.mock.ts     |  6 +-----
 functions/src/services/user/userService.ts    |  1 -
 6 files changed, 33 insertions(+), 39 deletions(-)

diff --git a/functions/models/src/types/userRegistration.ts b/functions/models/src/types/userRegistration.ts
index e57c265c..ae23cc55 100644
--- a/functions/models/src/types/userRegistration.ts
+++ b/functions/models/src/types/userRegistration.ts
@@ -60,8 +60,13 @@ export const userRegistrationConverter = new Lazy(
     }),
 )
 
+export interface UserClaims {
+  type: UserType
+  organization?: string
+}
+
 export class UserRegistration {
-  // Properties
+  // Stored Properties
 
   readonly type: UserType
   readonly organization?: string
@@ -80,6 +85,18 @@ export class UserRegistration {
   readonly language?: string
   readonly timeZone?: string
 
+  // Computed Properties
+
+  get claims(): UserClaims {
+    const result: UserClaims = {
+      type: this.type,
+    }
+    if (this.organization !== undefined) {
+      result.organization = this.organization
+    }
+    return result
+  }
+
   // Constructor
 
   constructor(input: {
diff --git a/functions/src/functions/blocking.ts b/functions/src/functions/blocking.ts
index 8af9fe21..57e05b2e 100644
--- a/functions/src/functions/blocking.ts
+++ b/functions/src/functions/blocking.ts
@@ -25,7 +25,7 @@ export const beforeUserCreatedFunction = beforeUserCreated(
     const credential = event.credential
 
     // Escape hatch for users using invitation code to enroll
-    if (!credential) return
+    if (!credential) return {}
 
     if (event.data.email === undefined)
       throw new https.HttpsError(
@@ -64,6 +64,8 @@ export const beforeUserCreatedFunction = beforeUserCreated(
       isSingleSignOn: true,
     })
     await factory.trigger().userEnrolled(userDoc)
+
+    return { customClaims: invitation.content.user.claims as object }
   },
 )
 
@@ -72,12 +74,16 @@ export const beforeUserSignedInFunction = beforeUserSignedIn(
   async (event) => {
     try {
       const userService = getServiceFactory().user()
-      const claims = await userService.getClaims(event.data.uid)
-      logger.info(`beforeUserSignedIn finished successfully.`)
-      return { sessionClaims: claims }
+      const user = await userService.getUser(event.data.uid)
+      if (user !== undefined) {
+        logger.info(`beforeUserSignedIn finished successfully.`)
+        return { customClaims: user.content.claims as object }
+      }
+      logger.info(`beforeUserSignedIn finished without user.`)
+      return { customClaims: {} }
     } catch (error) {
       logger.error(`beforeUserSignedIn finished with error: ${String(error)}`)
-      return { sessionClaims: {} }
+      return { customClaims: {} }
     }
   },
 )
diff --git a/functions/src/services/user/databaseUserService.test.ts b/functions/src/services/user/databaseUserService.test.ts
index 7aff6798..7923d717 100644
--- a/functions/src/services/user/databaseUserService.test.ts
+++ b/functions/src/services/user/databaseUserService.test.ts
@@ -67,9 +67,7 @@ describe('DatabaseUserService', () => {
       expect(userData).to.exist
       expect(userData?.invitationCode).to.equal(invitationCode)
       expect(userData?.dateOfEnrollment).to.exist
-
-      const claims = await userService.getClaims(userId)
-      expect(claims).to.deep.equal({
+      expect(userData?.claims).to.deep.equal({
         type: UserType.admin,
       })
     })
@@ -113,9 +111,7 @@ describe('DatabaseUserService', () => {
       expect(userData).to.exist
       expect(userData?.invitationCode).to.equal(invitationCode)
       expect(userData?.dateOfEnrollment).to.exist
-
-      const claims = await userService.getClaims(userId)
-      expect(claims).to.deep.equal({
+      expect(userData?.claims).to.deep.equal({
         type: UserType.clinician,
         organization: 'mockOrganization',
       })
@@ -162,9 +158,7 @@ describe('DatabaseUserService', () => {
       expect(userData).to.exist
       expect(userData?.invitationCode).to.equal(invitationCode)
       expect(userData?.dateOfEnrollment).to.exist
-
-      const claims = await userService.getClaims(userId)
-      expect(claims).to.deep.equal({
+      expect(userData?.claims).to.deep.equal({
         type: UserType.patient,
         organization: 'mockOrganization',
       })
diff --git a/functions/src/services/user/databaseUserService.ts b/functions/src/services/user/databaseUserService.ts
index fcaf3c74..7767f567 100644
--- a/functions/src/services/user/databaseUserService.ts
+++ b/functions/src/services/user/databaseUserService.ts
@@ -25,11 +25,6 @@ import {
   type DatabaseService,
 } from '../database/databaseService.js'
 
-export interface UserClaims {
-  type: UserType
-  organization?: string
-}
-
 export class DatabaseUserService implements UserService {
   // Properties
 
@@ -63,19 +58,6 @@ export class DatabaseUserService implements UserService {
     })
   }
 
-  async getClaims(userId: string): Promise<object> {
-    const user = await this.getUser(userId)
-    if (user !== undefined) {
-      const claims: UserClaims = {
-        type: user.content.type,
-      }
-      if (user.content.organization !== undefined)
-        claims.organization = user.content.organization
-      return claims
-    }
-    return {}
-  }
-
   // Invitations
 
   async createInvitation(content: Invitation): Promise<{ id: string }> {
diff --git a/functions/src/services/user/userService.mock.ts b/functions/src/services/user/userService.mock.ts
index f3b49942..f13c5309 100644
--- a/functions/src/services/user/userService.mock.ts
+++ b/functions/src/services/user/userService.mock.ts
@@ -43,11 +43,7 @@ export class MockUserService implements UserService {
   async updateAuth(userId: string, user: UserAuth): Promise<void> {
     return
   }
-
-  async getClaims(userId: string): Promise<object> {
-    return {}
-  }
-
+  
   // Methods - Invitations
 
   async createInvitation(content: Invitation): Promise<{ id: string }> {
diff --git a/functions/src/services/user/userService.ts b/functions/src/services/user/userService.ts
index 3b050209..91c638f2 100644
--- a/functions/src/services/user/userService.ts
+++ b/functions/src/services/user/userService.ts
@@ -24,7 +24,6 @@ export interface UserService {
 
   getAuth(userId: string): Promise<UserAuth>
   updateAuth(userId: string, auth: UserAuth): Promise<void>
-  getClaims(userId: string): Promise<object>
 
   // Invitations