feat: refactor fetch middleware and header forwarding

package.json

@@ -24,6 +24,7 @@
     "@angular/ssr": "^19.0.3",
     "@trpc/client": "11.0.0-rc.643",
     "@trpc/server": "11.0.0-rc.643",
+    "cookie": "^1.0.2",
     "cors": "^2.8.5",
     "express": "^4.21.2",
     "rxjs": "~7.8.1",

projects/ngx-trpc/ng-package.json

@@ -4,5 +4,5 @@
   "lib": {
     "entryFile": "src/public-api.ts"
   },
-  "allowedNonPeerDependencies": ["superjson", "@trpc/client", "@trpc/server"]
+  "allowedNonPeerDependencies": ["superjson", "@trpc/client", "@trpc/server", "cookie"]
 }

projects/ngx-trpc/package.json

@@ -18,7 +18,8 @@
     "tslib": "^2.3.0",
     "superjson": "^2.2.1",
     "@trpc/client": "^11.0.0-rc.643",
-    "@trpc/server": "^11.0.0-rc.643"
+    "@trpc/server": "^11.0.0-rc.643",
+    "cookie": "^1.0.2"
   },
   "sideEffects": false,
   "publishConfig": {

projects/ngx-trpc/src/lib/libs/mutex.util.ts

@@ -0,0 +1,42 @@
+export class Mutex {
+  private queue: (() => void)[] = [];
+  private locked = false;
+
+  async acquire(): Promise<() => void> {
+    return new Promise<() => void>((resolve) => {
+      const release = () => {
+        this.locked = false;
+        const next = this.queue.shift();
+        if (next) {
+          next();
+        }
+      };
+
+      if (!this.locked) {
+        this.locked = true;
+        resolve(release);
+      } else {
+        this.queue.push(() => {
+          this.locked = true;
+          resolve(release);
+        });
+      }
+    });
+  }
+}
+
+export async function wrapInMutex<T>(
+  fn: () => Promise<T>,
+  mutex: Mutex,
+  disable?: boolean
+): Promise<T> {
+  if (disable) {
+    return await fn();
+  }
+  const release = await mutex.acquire();
+  try {
+    return await fn();
+  } finally {
+    release();
+  }
+}

projects/ngx-trpc/src/lib/trpc.config.ts

@@ -38,5 +38,11 @@ export interface ITrpcConfig {
      * Default: `['set-cookie']`
      */
     forwardHeaders?: string[];
+
+    /**
+     * Disable sequential requests in SSR context.
+     * Only enable this, if you don't set cookies in createContext while doing sequential requests in ssr.
+     */
+    disableSequentialRequests?: boolean;
   };
 }

projects/ngx-trpc/src/lib/trpc.provider.ts

@@ -1,4 +1,11 @@
-import {InjectionToken, PLATFORM_ID, Provider, TransferState} from '@angular/core';
+import {
+  InjectionToken,
+  PLATFORM_ID,
+  Provider,
+  REQUEST,
+  RESPONSE_INIT,
+  TransferState
+} from '@angular/core';
 import {ITrpcConfig, provideTrpcConfig} from './trpc.config';
 import {AnyRouter} from '@trpc/server';
 import {CreateTRPCClient, createTRPCRxJSProxyClient} from './rxjs-proxy/create-rxjs-client';
@@ -10,8 +17,8 @@ import {
   tRPC_CACHE_STATE
 } from './utils/cache-state';
 import {transferStateLink} from './utils/transfer-state-link';
-import {FetchHttpClient} from './utils/fetch-http-client';
 import {getPlatformConfig, normalizeWebSocketUrl} from './utils/config-utils';
+import {FetchMiddleware} from './utils/fetch.middleware';
 
 export type TrpcClient<TRouter extends AnyRouter> = CreateTRPCClient<TRouter>;
 
@@ -29,14 +36,16 @@ export function provideTrpc<AppRouter extends AnyRouter>(
     provideTrpcCacheStateStatusManager(),
     {
       provide: token,
-      useFactory: (fetchHttpClient: FetchHttpClient, platformId: Object) => {
+      useFactory: (req: Request | null, res: ResponseInit | null, platformId: Object) => {
         const _isBrowser = isPlatformBrowser(platformId);
 
         const httpConfig = getPlatformConfig(_isBrowser, config.http, config.ssr?.http);
 
+        const fetchMiddleware = new FetchMiddleware(config, req, res);
+
         const trpcHttpLink = httpBatchLink({
           url: httpConfig.url,
-          fetch: fetchHttpClient.fetch.bind(fetchHttpClient)
+          fetch: _isBrowser ? undefined : (input, init) => fetchMiddleware.fetch(input, init)
         });
 
         let link: TRPCLink<AnyRouter> = trpcHttpLink;
@@ -57,7 +66,7 @@ export function provideTrpc<AppRouter extends AnyRouter>(
           links: [transferStateLink(), link]
         });
       },
-      deps: [FetchHttpClient, PLATFORM_ID, tRPC_CACHE_STATE, TransferState]
+      deps: [REQUEST, RESPONSE_INIT, PLATFORM_ID, tRPC_CACHE_STATE, TransferState]
     }
   ];
 }

projects/ngx-trpc/src/lib/utils/fetch.middleware.ts

@@ -0,0 +1,67 @@
+import {ITrpcConfig} from '../trpc.config';
+import * as cookie from 'cookie';
+import {Mutex, wrapInMutex} from '../libs/mutex.util';
+
+export class FetchMiddleware {
+  private _setCookiesCache?: string;
+
+  private _mutex = new Mutex();
+
+  constructor(
+    private _config: ITrpcConfig,
+    private _request: Request | null,
+    private _response: ResponseInit | null
+  ) {}
+
+  async fetch(input: RequestInfo | URL | string, init?: RequestInit) {
+    // Wrap this in mutex. Because of setCookieCache, we need to make sequential requests in SSR.
+    // This should not be a problem, since we use batch calls.
+    return wrapInMutex(
+      async () => {
+        if (typeof input != 'string') {
+          throw new Error('[ngx-trpc] Only string urls are supported right now.');
+        }
+
+        if (this._request) {
+          const headers: HeadersInit = {};
+          this._request.headers.forEach((value: string, key: string) => {
+            headers[key] = value;
+          });
+
+          if (this._setCookiesCache) {
+            const cookies = cookie.parse(headers['cookie'] || '');
+            const newCookies = cookie.parse(this._setCookiesCache);
+
+            for (let key in newCookies) {
+              cookies[key] = newCookies[key];
+            }
+
+            headers['cookie'] = Object.entries(cookies)
+              .filter(([_, value]) => value !== undefined)
+              .map(([key, value]) => cookie.serialize(key, value!))
+              .join('; ');
+          }
+
+          init = {...init, headers};
+        }
+
+        const r = await fetch(input, init);
+        if (this._response && this._response.headers && this._response.headers instanceof Headers) {
+          for (let [key, value] of r.headers) {
+            if (this._config.ssr?.forwardHeaders?.includes(key) || key === 'set-cookie') {
+              this._response.headers.set(key, value);
+            }
+
+            // set the cookie in the angular request object so we can use it in sequential requests when not using batchLink
+            if (key == 'set-cookie') {
+              this._setCookiesCache = value;
+            }
+          }
+        }
+        return r;
+      },
+      this._mutex,
+      this._config.ssr?.disableSequentialRequests
+    );
+  }
+}