node.jsnestjsnestjs-passportasync-hooks

AsyncLocalStorage Works in NestJS Guard's canActivate method but Not in Passport Strategy's validate Method


I'm working on a NestJS application where I need to maintain request-specific context using AsyncLocalStorage. I have a ContextRequestService that uses AsyncLocalStorage to store and retrieve context data throughout the request lifecycle.

Here's my ContextRequestService:

import { Injectable } from '@nestjs/common';
import { AsyncLocalStorage } from 'async_hooks';

interface ContextStore {
  lang: string;
  requestId: string;
  userId?: number;
  clinicId?: number;
}

@Injectable()
export class ContextRequestService {
  constructor(
    private readonly asyncLocalStorage: AsyncLocalStorage<ContextStore>,
  ) {}

  run(data: ContextStore, callback: () => void) {
    this.asyncLocalStorage.run(data, callback);
  }

  getLang(): string {
    const store = this.asyncLocalStorage.getStore();
    return store?.lang || 'en';
  }

  getRequestId(): string {
    const store = this.asyncLocalStorage.getStore();
    return store?.requestId || 'N/A';
  }

  setUserId(id: number) {
    const store = this.asyncLocalStorage.getStore();
    this.asyncLocalStorage.enterWith({ ...store, userId: id });
  }

  getUserId() {
    const store = this.asyncLocalStorage.getStore();
    return store?.userId;
  }

  setClinicId(id: number) {
    const store = this.asyncLocalStorage.getStore();
    this.asyncLocalStorage.enterWith({ ...store, clinicId: id });
  }

  getClinicId() {
    const store = this.asyncLocalStorage.getStore();
    return store?.clinicId;
  }
}

the first mode:

In my custom JwtAuthGuard, I'm able to set and retrieve data from the AsyncLocalStorage without any issues:

import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { AuthGuard } from "@nestjs/passport";
import { SetMetadata } from '@nestjs/common';
import { Reflector } from '@nestjs/core';
import { ContextRequestService } from "../common/context/context-request";

export const IS_PUBLIC_KEY = 'isPublic';
export const Public = () => SetMetadata(IS_PUBLIC_KEY, true);
@Injectable()
export class JwtAuthGuard extends AuthGuard('jwt') implements CanActivate {
  constructor(private reflector: Reflector, private readonly contextRequestService: ContextRequestService) {
    super();
  }

  canActivate(context: ExecutionContext) {
    const isPublic = this.reflector.getAllAndOverride<boolean>(IS_PUBLIC_KEY, [
      context.getHandler(),
      context.getClass(),
    ]);
    if (isPublic) {
      return true;
    }
    this.contextRequestService.setClinicId(10)
    return super.canActivate(context);
  }
}

import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { ContextRequestService } from '../../common/context/context-request';
import { AsyncLocalStorage } from 'async_hooks';

interface ContextStore {
  lang: string;
  requestId: string;
  userId?: number;
  clinicId?: number;
}
@Injectable()
export class AccessGuard implements CanActivate {
  constructor(
    private readonly contextRequest: ContextRequestService,
    private readonly asyncLocalStorage: AsyncLocalStorage<ContextStore>,
    ) {}

  async canActivate(context: ExecutionContext): Promise<boolean> {

    ...
    
    console.log("____first mode_____", this.asyncLocalStorage.getStore());
    
    ...
  }
 }

the second mode:

import { PassportStrategy } from "@nestjs/passport";
import { ExtractJwt, Strategy } from "passport-jwt";
import { ContextRequestService } from "../common/context/context-request";
import {Injectable} from '@nestjs/common';
import { AsyncLocalStorage } from "async_hooks";

interface ContextStore {
  lang: string;
  requestId: string;
  userId?: number;
  clinicId?: number;
}
@Injectable()
export class JwtStrategy extends PassportStrategy(Strategy) {
  constructor( private readonly contextRequestService: ContextRequestService,
    private readonly asyncLocalStorage: AsyncLocalStorage<ContextStore>) {
    super({
      jwtFromRequest: ExtractJwt.fromAuthHeaderAsBearerToken(),
      ignoreExpiration: false,
      secretOrKey: process.env.JWT_SECRET,
    });
  }
  validate = async (payload: { userId: string, clinicId: number }) => {
    this.contextRequestService.setUserId(parseInt(payload.userId));
    this.contextRequestService.setClinicId(1000000);
    console.log("______ the second mode__ JWT Strategy__", this.asyncLocalStorage.getStore());
    return { userId: payload.userId };
  }
}

import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { ContextRequestService } from '../../common/context/context-request';
import { AsyncLocalStorage } from 'async_hooks';

interface ContextStore {
  lang: string;
  requestId: string;
  userId?: number;
  clinicId?: number;
}
@Injectable()
export class AccessGuard implements CanActivate {
  constructor(
    private readonly asyncLocalStorage: AsyncLocalStorage<ContextStore>,
    ) {}

  async canActivate(context: ExecutionContext): Promise<boolean> {

    ...
    
    console.log("____second mode_____AccessGuard ", this.asyncLocalStorage.getStore());
    ...

  }
 }

Question:

Why does AsyncLocalStorage work in the canActivate method of my guard but not in the validate method of my Passport strategy?

output first mode :

____first mode_____ {
  lang: 'en',
  requestId: '0e7fb339-7850-410f-a46a-d40e51bbb52e',
  clinicId: 10
}

output the scond mode:

______ the second mode__ JWT Strategy__ {
  lang: 'en',
  requestId: '2593cef0-407d-42bc-9d88-dfb1c8d657c5',
  clinicId: 1000000,
  userId: 1
}
____second mode_____AccessGuard  {
  lang: 'en',
  requestId: '2593cef0-407d-42bc-9d88-dfb1c8d657c5'
}


Solution

  • Short answer, don't use .enterWith and spread operators, unless you're ABSOLUTELY SURE that you (and all of the node modules you're using) aren't doing .then anywhere during the context of that request. These two things by design aren't always compatible.


    Long answer (oversimplified):

    If you're using AsyncLocalStorage, every time you create a Promise, AsyncLocalStorage creates a new pointer to the storage and gives your promise THAT. Afterward, it takes the current value of that pointer and puts it back on the storage.

    This isn't the code, but it draws the picture:

    let globalStore = { requestId: 1 }
    ;() => {
      let innerStore = globalStore
      await doPromise()
      globalStore = innerStore
    })()
    

    and when you do .enterWith, it replaces the storage pointer with the value you give it:

    innerStore = { ...innerStore } // and spread overwrites the pointer, passing by values
    

    Now suddenly this code:

    const store = this.asyncLocalStorage.getStore();
    this.asyncLocalStorage.enterWith({ ...store, userId: id });
    

    is no longer the same as this code:

    const store = this.asyncLocalStorage.getStore();
    store.userId = id;
    

    If any code ANYWHERE uses .then, you have created TWO TWO promises at the same time, and you know what it does when it creates a promise... Here it is with .enterWith:

    let globalStore = { requestId: 1 }
    // let promise1 = first()
    let innerStore1 = globalStore
    // promise2 = promise1.then(second)
    let innerStore2 = globalStore
    // executing first
    innerStore1 = { ...innerStore1, key1: value1 } // override the reference with values
    // executing second
    innerStore2 = { ...innerStore2, key2: value2 } // override the reference with values
    // await promise1
    globalStore = innerStore1 // values from innerStore1, which does not include key2
    // await promise2
    globalStore = innerStore2 // values from innerStore2, which does not include key1
    

    but if you set properties on the singleton store instead of giving it a new object using .enterWith:

    let globalStore = { requestId: 1 }
    // let promise1 = first()
    let innerStore1 = globalStore
    // promise2 = promise1.then(second)
    let innerStore2 = globalStore
    // executing first
    innerStore1.key1 = value1 // updated the reference
    // executing second
    innerStore2.key2 = value2 // updated the reference
    // await promise1
    globalStore = innerStore1 // same reference
    // await promise2
    globalStore = innerStore2 // same reference
    

    and if you use .enterWith, but await instead of .then to chain your promises:

    let globalStore = { requestId: 1 }
    
    // await first()
    let innerStore1 = globalStore
    innerStore1 = { ...innerStore1, key1: value1 }
    globalStore = innerStore1 // includes key1
    
    // await second()
    let innerStore2 = globalStore // includes key1
    innerStore2 = { ...innerStore2, key2: value2 }
    globalStore = innerStore2 // includes key1 & key2
    

    So here is your fixed ContextRequestService:

    export class ContextRequestService {
      constructor(
        private readonly asyncLocalStorage: AsyncLocalStorage<ContextStore>,
      ) {}
    
      run(data: ContextStore, callback: () => void) {
        this.asyncLocalStorage.run(data, callback);
      }
    
      getLang(): string {
        const store = this.asyncLocalStorage.getStore();
        return store?.lang || 'en';
      }
    
      getRequestId(): string {
        const store = this.asyncLocalStorage.getStore();
        return store?.requestId || 'N/A';
      }
    
      setUserId(id: number) {
        const store = this.asyncLocalStorage.getStore();
        store.userId = id;
      }
    
      getUserId() {
        const store = this.asyncLocalStorage.getStore();
        return store?.userId;
      }
    
      setClinicId(id: number) {
        const store = this.asyncLocalStorage.getStore();
        store.clinicId = id;
      }
    
      getClinicId() {
        const store = this.asyncLocalStorage.getStore();
        return store?.clinicId;
      }
    }