fix: add explicit schema for token payload

example/app.ts

@@ -21,10 +21,15 @@ export default await AppContainer.create({
               ability.can('manage', 'all');
             }
           },
-          loginCredentialsSchema: z.object({
-            password: z.string().min(1),
-            username: z.string().min(1)
-          }),
+          schemas: {
+            loginCredentials: z.object({
+              password: z.string().min(1),
+              username: z.string().min(1)
+            }),
+            tokenPayload: z.object({
+              isAdmin: z.boolean()
+            })
+          },
           userQuery: async ({ username }) => {
             if (username !== 'admin') {
               return null;

src/index.ts

@@ -3,7 +3,7 @@ export { CurrentUser } from './decorators/current-user.decorator.js';
 export { RouteAccess } from './decorators/route-access.decorator.js';
 export { ValidationSchema } from './decorators/validation-schema.decorator.js';
 export { DataTransferObject } from './mixins/data-transfer-object.mixin.js';
-export type { AuthModuleOptions, UserQuery, UserQueryResult } from './modules/auth/auth.config.js';
+export type { AuthModuleOptions } from './modules/auth/auth.config.js';
 export { AuthModule } from './modules/auth/auth.module.js';
 export { ConfigService } from './modules/config/config.service.js';
 export { CryptoService } from './modules/crypto/crypto.service.js';

src/modules/auth/__tests__/auth.module.spec.ts

@@ -43,7 +43,7 @@ describe('AuthModule', () => {
   let jwtStrategy: JwtStrategy;
 
   let defineAbility: Mock<DefineAbility>;
-  let userQuery: Mock<UserQuery>;
+  let userQuery: Mock<UserQuery<{ password: string; username: string }, { username: string }>>;
 
   const loginCredentials = {
     password: 'password',
@@ -62,10 +62,15 @@ describe('AuthModule', () => {
       imports: [
         AuthModule.forRoot({
           defineAbility,
-          loginCredentialsSchema: z.object({
-            password: z.string(),
-            username: z.string()
-          }),
+          schemas: {
+            loginCredentials: z.object({
+              password: z.string(),
+              username: z.string()
+            }),
+            tokenPayload: z.object({
+              username: z.string()
+            })
+          },
           userQuery: userQuery
         }),
         ConfigModule.forRoot({

src/modules/auth/auth.config.ts

@@ -2,7 +2,7 @@
 import type { RawRuleOf } from '@casl/ability';
 import type { PrismaQuery, Subjects } from '@casl/prisma';
 import { ConfigurableModuleBuilder } from '@nestjs/common';
 import { Prisma } from '@prisma/client';
 import type { DefaultSelection } from '@prisma/client/runtime/library';
 import type { z } from 'zod';
 
@@ -47,41 +47,37 @@
 
 type BaseLoginCredentialsSchema = z.ZodType<BaseLoginCredentials>;
 
-type UserQueryResult = {
-  hashedPassword: string;
-  tokenPayload: {
-    [key: string]: unknown;
-  };
-};
-
 interface JwtPayload {
   [key: string]: any;
   permissions: Permission[];
 }
 
-type UserQuery<TLoginCredentials extends BaseLoginCredentials = BaseLoginCredentials> = (
-  credentials: TLoginCredentials
-) => Promise<null | UserQueryResult>;
+type UserQuery<
+  TLoginCredentials extends BaseLoginCredentials = BaseLoginCredentials,
+  TPayload extends { [key: string]: unknown } = { [key: string]: unknown }
+> = (credentials: TLoginCredentials) => Promise<null | {
+  hashedPassword: string;
+  tokenPayload: TPayload;
+}>;
 
 type LoginResponseBody = {
   accessToken: string;
 };
 
 type AuthModuleOptions<
   TLoginCredentialsSchema extends BaseLoginCredentialsSchema = BaseLoginCredentialsSchema,
-  TUserQuery extends UserQuery<z.TypeOf<TLoginCredentialsSchema>> = UserQuery<z.TypeOf<TLoginCredentialsSchema>>
+  TPayloadSchema extends z.ZodType<{ [key: string]: unknown }> = z.ZodType<{ [key: string]: unknown }>
 > = {
-  defineAbility: TUserQuery extends (...args: any[]) => Promise<infer R>
-    ? R extends { tokenPayload: infer TPayload extends { [key: string]: unknown } }
-      ? DefineAbility<TPayload>
-      : never
-    : never;
-  loginCredentialsSchema: TLoginCredentialsSchema;
-  userQuery: TUserQuery;
+  defineAbility: (ability: AbilityBuilder<AppAbility>, tokenPayload: z.TypeOf<TPayloadSchema>) => any;
+  schemas: {
+    loginCredentials: TLoginCredentialsSchema;
+    tokenPayload: TPayloadSchema;
+  };
+  userQuery: UserQuery<z.TypeOf<TLoginCredentialsSchema>, z.TypeOf<TPayloadSchema>>;
 };
 
 export const { ConfigurableModuleClass: ConfigurableAuthModule, MODULE_OPTIONS_TOKEN: AUTH_MODULE_OPTIONS_TOKEN } =
-  new ConfigurableModuleBuilder<AuthModuleOptions<any>>()
+  new ConfigurableModuleBuilder<AuthModuleOptions<any, any>>()
     .setClassMethodName('forRoot')
     .setExtras({}, (definition) => ({
       ...definition,
@@ -107,6 +103,5 @@
   JwtPayload,
   LoginResponseBody,
   Permission,
-  UserQuery,
-  UserQueryResult
+  UserQuery
 };

src/modules/auth/auth.module.ts

@@ -19,7 +19,7 @@ import { LoginCredentialsDto } from './dto/login-credentials.dto.js';
 import { JwtAuthGuard } from './guards/jwt-auth.guard.js';
 import { JwtStrategy } from './strategies/jwt.strategy.js';
 
-import type { AuthModuleOptions, BaseLoginCredentialsSchema, UserQuery } from './auth.config.js';
+import type { AuthModuleOptions, BaseLoginCredentialsSchema } from './auth.config.js';
 
 @Module({
   controllers: [AuthController],
@@ -58,22 +58,22 @@ import type { AuthModuleOptions, BaseLoginCredentialsSchema, UserQuery } from '.
 export class AuthModule extends ConfigurableAuthModule implements OnModuleInit {
   private readonly loginCredentialsSchema: BaseLoginCredentialsSchema;
 
-  constructor(@Inject(AUTH_MODULE_OPTIONS_TOKEN) { loginCredentialsSchema }: AuthModuleOptions) {
+  constructor(@Inject(AUTH_MODULE_OPTIONS_TOKEN) { schemas }: AuthModuleOptions) {
     super();
-    this.loginCredentialsSchema = loginCredentialsSchema;
+    this.loginCredentialsSchema = schemas.loginCredentials;
   }
 
   static forRoot<
     TLoginCredentialsSchema extends BaseLoginCredentialsSchema,
-    TUserQuery extends UserQuery<z.TypeOf<TLoginCredentialsSchema>>
-  >(options: AuthModuleOptions<TLoginCredentialsSchema, TUserQuery>): DynamicModule {
+    TPayloadSchema extends z.ZodType<{ [key: string]: unknown }>
+  >(options: AuthModuleOptions<TLoginCredentialsSchema, TPayloadSchema>): DynamicModule {
     return super.forRoot(options);
   }
   static forRootAsync<
     TLoginCredentialsSchema extends BaseLoginCredentialsSchema,
-    TUserQuery extends UserQuery<z.TypeOf<TLoginCredentialsSchema>>
+    TPayloadSchema extends z.ZodType<{ [key: string]: unknown }>
   >(
-    options: ConfigurableModuleAsyncOptions<AuthModuleOptions<TLoginCredentialsSchema, TUserQuery>, 'create'>
+    options: ConfigurableModuleAsyncOptions<AuthModuleOptions<TLoginCredentialsSchema, TPayloadSchema>, 'create'>
   ): DynamicModule {
     return super.forRootAsync(options);
   }