From 7c701ded5f205b1fbfd7c2d37ecc2a594ad5c387 Mon Sep 17 00:00:00 2001
From: Achintha Isuru <achinthaisuru444@gmail.com>
Date: Sat, 28 Feb 2026 18:17:49 -0500
Subject: [PATCH] feat: Enhance authentication by refining user role validation
 during session handling and ensuring immediate ID token refresh after sign-in
 to prevent unauthenticated Data Connect SDK requests.

---
 .../mixins/session_handler_mixin.dart         | 22 ++++++++++++++-----
 .../auth_repository_impl.dart                 | 10 +++++++++
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/apps/mobile/packages/data_connect/lib/src/services/mixins/session_handler_mixin.dart b/apps/mobile/packages/data_connect/lib/src/services/mixins/session_handler_mixin.dart
index d04a2cb3..0ce10c6a 100644
--- a/apps/mobile/packages/data_connect/lib/src/services/mixins/session_handler_mixin.dart
+++ b/apps/mobile/packages/data_connect/lib/src/services/mixins/session_handler_mixin.dart
@@ -205,13 +205,23 @@ mixin SessionHandlerMixin {
     try {
       _emitSessionState(SessionState.loading());
 
-      // Validate role if allowed roles are specified
+      // Validate role only when allowed roles are specified.
       if (_allowedRoles.isNotEmpty) {
-        final bool isAuthorized = await validateUserRole(
-          user.uid,
-          _allowedRoles,
-        );
-        if (!isAuthorized) {
+        final String? userRole = await fetchUserRole(user.uid);
+
+        if (userRole == null) {
+          // User has no record in the database yet. This is expected during
+          // the sign-up flow: Firebase Auth fires authStateChanges before the
+          // repository has created the PostgreSQL user record. Do NOT sign out —
+          // just emit unauthenticated and let the registration flow complete.
+          _emitSessionState(SessionState.unauthenticated());
+          return;
+        }
+
+        if (!_allowedRoles.contains(userRole)) {
+          // User IS in the database but has a role that is not permitted in
+          // this app (e.g., a STAFF-only user trying to use the Client app).
+          // Sign them out to force them to use the correct app.
           await auth.signOut();
           _emitSessionState(SessionState.unauthenticated());
           return;
diff --git a/apps/mobile/packages/features/client/authentication/lib/src/data/repositories_impl/auth_repository_impl.dart b/apps/mobile/packages/features/client/authentication/lib/src/data/repositories_impl/auth_repository_impl.dart
index f3a195fe..d21ac6ce 100644
--- a/apps/mobile/packages/features/client/authentication/lib/src/data/repositories_impl/auth_repository_impl.dart
+++ b/apps/mobile/packages/features/client/authentication/lib/src/data/repositories_impl/auth_repository_impl.dart
@@ -90,6 +90,12 @@ class AuthRepositoryImpl implements AuthRepositoryInterface {
         );
       }
 
+      // Force-refresh the ID token so the Data Connect SDK has a valid bearer
+      // token before we fire any mutations. Without this, there is a race
+      // condition where the gRPC layer sends the request unauthenticated
+      // immediately after account creation (gRPC code 16 UNAUTHENTICATED).
+      await firebaseUser.getIdToken(true);
+
       // New user created successfully, proceed to create PostgreSQL entities
       return await _createBusinessAndUser(
         firebaseUser: firebaseUser,
@@ -165,6 +171,10 @@ class AuthRepositoryImpl implements AuthRepositoryInterface {
         );
       }
 
+      // Force-refresh the ID token so the Data Connect SDK receives a valid
+      // bearer token before any subsequent Data Connect queries run.
+      await firebaseUser.getIdToken(true);
+
       // Sign-in succeeded! Check if user already has a BUSINESS account in PostgreSQL
       final bool hasBusinessAccount = await _checkBusinessUserExists(
         firebaseUser.uid,