From 45ccd395622db4fe45cdd812f18b56b098dc3d89 Mon Sep 17 00:00:00 2001
From: James Daniels <jamesdaniels@google.com>
Date: Fri, 13 Dec 2024 16:35:27 -0500
Subject: [PATCH] Adding more injector safety (#3590)

* Zone wrapper noops for our other helpers
* Add a warning / error on potential Zone / hydration issues
* Pass injection context to `zoneWrapFn`
* Pass injection context into the Promise wrapper
* `beforeAuthStateChanged` should not block
---
 src/auth/firebase.ts |  2 +-
 src/zones.ts         | 63 +++++++++++++++++++++++++++++++++++++-------
 tools/build.ts       |  1 +
 3 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/src/auth/firebase.ts b/src/auth/firebase.ts
index d1a25919a..97c07f167 100644
--- a/src/auth/firebase.ts
+++ b/src/auth/firebase.ts
@@ -59,7 +59,7 @@ import {
 } from 'firebase/auth';
 
 export const applyActionCode = ɵzoneWrap(_applyActionCode, true);
-export const beforeAuthStateChanged = ɵzoneWrap(_beforeAuthStateChanged, true);
+export const beforeAuthStateChanged = ɵzoneWrap(_beforeAuthStateChanged, false);
 export const checkActionCode = ɵzoneWrap(_checkActionCode, true);
 export const confirmPasswordReset = ɵzoneWrap(_confirmPasswordReset, true);
 export const connectAuthEmulator = ɵzoneWrap(_connectAuthEmulator, true);
diff --git a/src/zones.ts b/src/zones.ts
index c32ac1dd5..d8505433e 100644
--- a/src/zones.ts
+++ b/src/zones.ts
@@ -4,7 +4,9 @@ import {
   Injector,
   NgZone,
   PendingTasks,
-  inject
+  inject,
+  isDevMode,
+  runInInjectionContext
 } from '@angular/core';
 import { pendingUntilEvent } from '@angular/core/rxjs-interop';
 import {
@@ -76,31 +78,71 @@ function getSchedulers() {
   return inject(ɵAngularFireSchedulers);
 }
 
+var alreadyWarned = false;
+function warnOutsideInjectionContext(original: any, operation: string) {
+  if (isDevMode()) {
+    console.warn(`Firebase API called outside injection context: ${operation}(${original.name})`);
+    if (!alreadyWarned) {
+      alreadyWarned = true;
+      console.error("Calling Firebase APIs outside of an Injection context may destabilize your application leading to subtle change-detection and hydration bugs. Find more at https://github.com/angular/angularfire/blob/main/docs/zones.md");
+    }
+  }
+}
+
 function runOutsideAngular<T>(fn: (...args: any[]) => T): T {
-  return inject(NgZone).runOutsideAngular(() => fn());
+  let ngZone: NgZone|undefined;
+  try {
+    ngZone = inject(NgZone);
+  } catch(e) {
+    warnOutsideInjectionContext(fn, "runOutsideAngular");
+  }
+  if (!ngZone) {return fn();}
+  return ngZone.runOutsideAngular(() => fn());
 }
 
 function run<T>(fn: (...args: any[]) => T): T {
-  return inject(NgZone).run(() => fn());
+  let ngZone: NgZone|undefined;
+  try {
+    ngZone = inject(NgZone);
+  } catch(e) {
+    warnOutsideInjectionContext(fn, "run");
+  }
+  if (!ngZone) {return fn();}
+  return ngZone.run(() => fn());
 }
 
 export function observeOutsideAngular<T>(obs$: Observable<T>): Observable<T> {
-  return obs$.pipe(observeOn(getSchedulers().outsideAngular));
+  let schedulers: ɵAngularFireSchedulers|undefined;
+  try {
+    schedulers = getSchedulers();
+  } catch(e) {
+    warnOutsideInjectionContext(obs$, "observeOutsideAngular");
+  }
+  if (!schedulers) {return obs$;}
+  return obs$.pipe(observeOn(schedulers.outsideAngular));
 }
 
 export function observeInsideAngular<T>(obs$: Observable<T>): Observable<T> {
-  return obs$.pipe(observeOn(getSchedulers().insideAngular));
+  let schedulers: ɵAngularFireSchedulers|undefined;
+  try {
+    schedulers = getSchedulers();
+  } catch(e) {
+    warnOutsideInjectionContext(obs$, "observeInsideAngular");
+  }
+  if (!schedulers) {return obs$;}
+  return obs$.pipe(observeOn(schedulers.insideAngular));
 }
 
 const zoneWrapFn = (
   it: (...args: any[]) => any,
-  taskDone: VoidFunction | undefined
+  taskDone: VoidFunction | undefined,
+  injector: Injector,
 ) => {
   return (...args: any[]) => {
     if (taskDone) {
       setTimeout(taskDone, 0);
     }
-    return run(() => it.apply(this, args));
+    return runInInjectionContext(injector, () => run(() => it.apply(this, args)));
   };
 };
 
@@ -117,6 +159,7 @@ export const ɵzoneWrap = <T= unknown>(it: T, blockUntilFirst: boolean): T => {
       pendingTasks = inject(PendingTasks);
       injector = inject(Injector);
     } catch(e) {
+      warnOutsideInjectionContext(it, "ɵzoneWrap");
       return (it as any).apply(this, _arguments);
     }
     // if this is a callback function, e.g, onSnapshot, we should create a pending task and complete it
@@ -127,7 +170,7 @@ export const ɵzoneWrap = <T= unknown>(it: T, blockUntilFirst: boolean): T => {
           taskDone ||= run(() => pendingTasks.add());
         }
         // TODO create a microtask to track callback functions
-        _arguments[i] = zoneWrapFn(_arguments[i], taskDone);
+        _arguments[i] = zoneWrapFn(_arguments[i], taskDone, injector);
       }
     }
     const ret = runOutsideAngular(() => (it as any).apply(this, _arguments));
@@ -153,8 +196,8 @@ export const ɵzoneWrap = <T= unknown>(it: T, blockUntilFirst: boolean): T => {
         () =>
           new Promise((resolve, reject) => {
             pendingTasks.run(() => ret).then(
-              (it) => run(() => resolve(it)),
-              (reason) => run(() => reject(reason))
+              (it) => runInInjectionContext(injector, () => run(() => resolve(it))),
+              (reason) => runInInjectionContext(injector, () => run(() => reject(reason)))
             );
           })
       );
diff --git a/tools/build.ts b/tools/build.ts
index 7da9a45b4..397d78906 100644
--- a/tools/build.ts
+++ b/tools/build.ts
@@ -75,6 +75,7 @@ ${exportedZoneWrappedFns}
       browserSessionPersistence: null,
       indexedDBLocalPersistence: null,
       prodErrorMap: null,
+      beforeAuthStateChanged: { blockUntilFirst: false },
     }),
     reexport('database', 'rxfire', 'rxfire/database', tsKeys<typeof import('rxfire/database')>()),
     reexport('database', 'firebase', 'firebase/database', tsKeys<typeof import('firebase/database')>()),