Angular - Enhancing MemoryTokenStorageService to support session storage for specific keys

npm/ng-packs/packages/oauth/src/lib/services/memory-token-storage.service.ts

@@ -22,6 +22,8 @@ export class MemoryTokenStorageService implements OAuthStorage {
     'granted_scopes',
   ];
 
+  private keysShouldStoreInSession = ['nonce', 'PKCE_verifier', 'session_state'];
+
   private worker?: any;
   private port?: MessagePort;
   private cache = new Map<string, string>();
@@ -89,38 +91,55 @@ export class MemoryTokenStorageService implements OAuthStorage {
   }
 
   getItem(key: string): string | null {
-    if (!this.keysShouldStoreInMemory.includes(key)) {
-      return this.localStorageService.getItem(key);
+    if (this.keysShouldStoreInMemory.includes(key)) {
+      return this.cache.get(key) || null;
+    }
+
+    if (this.keysShouldStoreInSession.includes(key)) {
+      return this.getSessionItem(key);
     }
-    return this.cache.get(key) || null;
+
+    return this.localStorageService.getItem(key);
   }
 
   setItem(key: string, value: string): void {
-    if (!this.keysShouldStoreInMemory.includes(key)) {
-      this.localStorageService.setItem(key, value);
+    if (this.keysShouldStoreInMemory.includes(key)) {
+      if (this.useSharedWorker && this.port) {
+        this.cache.set(key, value);
+        this.port.postMessage({ action: 'set', key, value });
+      } else {
+        this.cache.set(key, value);
+      }
+
       return;
     }
 
-    if (this.useSharedWorker && this.port) {
-      this.cache.set(key, value);
-      this.port.postMessage({ action: 'set', key, value });
-    } else {
-      this.cache.set(key, value);
+    if (this.keysShouldStoreInSession.includes(key)) {
+      this.setSessionItem(key, value);
+      return;
     }
+
+    this.localStorageService.setItem(key, value);
   }
 
   removeItem(key: string): void {
-    if (!this.keysShouldStoreInMemory.includes(key)) {
-      this.localStorageService.removeItem(key);
+    if (this.keysShouldStoreInMemory.includes(key)) {
+      if (this.useSharedWorker && this.port) {
+        this.cache.delete(key);
+        this.port.postMessage({ action: 'remove', key });
+      } else {
+        this.cache.delete(key);
+      }
+
       return;
     }
 
-    if (this.useSharedWorker && this.port) {
-      this.cache.delete(key);
-      this.port.postMessage({ action: 'remove', key });
-    } else {
-      this.cache.delete(key);
+    if (this.keysShouldStoreInSession.includes(key)) {
+      this.removeSessionItem(key);
+      return;
     }
+
+    this.localStorageService.removeItem(key);
   }
 
   clear(): void {
@@ -165,6 +184,37 @@ export class MemoryTokenStorageService implements OAuthStorage {
     }, 100);
   }
 
+  private getSessionStorage(): Storage | null {
+    try {
+      return this._document.defaultView?.sessionStorage ?? null;
+    } catch {
+      return null;
+    }
+  }
+
+  private getSessionItem(key: string): string | null {
+    const storage = this.getSessionStorage();
+    return storage ? storage.getItem(key) : this.localStorageService.getItem(key);
+  }
+
+  private setSessionItem(key: string, value: string): void {
+    const storage = this.getSessionStorage();
+    if (storage) {
+      storage.setItem(key, value);
+    } else {
+      this.localStorageService.setItem(key, value);
+    }
+  }
+
+  private removeSessionItem(key: string): void {
+    const storage = this.getSessionStorage();
+    if (storage) {
+      storage.removeItem(key);
+    } else {
+      this.localStorageService.removeItem(key);
+    }
+  }
+
   private createWorkerDataUrl(): string {
     const workerScript = `const tokenStore = new Map();
 const ports = new Set();