feat: custom auth works

Author: Ansonhkg

packages/auth-services/src/auth-server/src/createAuthServer.ts

@@ -10,10 +10,7 @@ import { mainAppQueue } from '../../queue-manager/src/bullmqSetup'; // Adjusted
 import { apiKeyGateAndTracking } from '../middleware/apiKeyGate'; // Adjusted path
 import { rateLimiter } from '../middleware/rateLimiter'; // Adjusted path
 import { resp } from './response-helpers/response-helpers'; // Adjusted path
-import {
-  customAuthLoginRoute,
-  customAuthVerifyTokenRoute,
-} from './routes/auth/custom-auth-routes';
+
 import { stytchEmailRoutes } from './routes/auth/stytch/stytch-email';
 import { stytchWhatsAppRoutes } from './routes/auth/stytch/stytch-otp';
 import { stytchSmsRoutes } from './routes/auth/stytch/stytch-sms';
@@ -129,10 +126,6 @@ export const createLitAuthServer = (
       stytchWhatsAppRoutes(groupApp);
       stytchTotpRoutes(groupApp);
 
-      // Custom Auth
-      customAuthLoginRoute(groupApp);
-      customAuthVerifyTokenRoute(groupApp);
-
       return groupApp;
     })
 

packages/auth-services/src/auth-server/src/response-helpers/response-helpers.ts

@@ -33,9 +33,18 @@ const _res200 = (data: any) => {
     status: 200,
   });
 };
+
+const _res401 = (error: string) => {
+  return new Response(BigIntStringify({ error }), {
+    headers: HEADERS,
+    status: 401,
+  });
+};
+
 export const resp: Record<string, (...args: any[]) => Response> = {
   SUCCESS: _res200,
   QUEUED: _res202,
   ERROR: _res500,
   BAD_REQUEST: _res400,
+  UNAUTHORIZED: _res401,
 };

packages/auth-services/src/auth-server/src/routes/pkp/mint.ts

@@ -0,0 +1,24 @@
+import { MintRequestRaw } from '@lit-protocol/networks/src/networks/vNaga/LitChainClient/schemas/MintRequestSchema';
+import { ElysiaInstance } from '../../types/ElysiaInstance.type';
+import { addJob } from '../../../../queue-manager/src/bullmqSetup';
+import { resp } from '../../response-helpers/response-helpers';
+import { mintPkpDoc } from '../../../../queue-manager/src/handlers/pkpMint/pkpMint.doc';
+
+export const mint = (app: ElysiaInstance) => {
+  app.post(
+    '/mint',
+    async ({ body }: { body: MintRequestRaw }) => {
+      try {
+        const job = await addJob('pkpMint', { requestBody: body });
+        return resp.QUEUED(job.id, 'PKP minting request queued successfully.');
+      } catch (error: any) {
+        console.error(`[API] Failed to add job 'pkpMint' to queue:`, error);
+        return resp.ERROR(
+          'Failed to queue PKP minting request.' + error.message
+        );
+      }
+    },
+    mintPkpDoc
+  );
+  return app;
+};

packages/auth-services/src/auth-server/src/routes/status.ts

@@ -0,0 +1,33 @@
+/**
+ * Job Status routes
+ * Handles status checking for queued jobs
+ */
+import { getJobStatus } from '../../../queue-manager/src/bullmqSetup'; // Adjusted path
+import { getStatusDoc } from '../../../queue-manager/src/handlers/status/getStatus.doc'; // Adjusted path
+import { resp } from '../response-helpers/response-helpers';
+import { ElysiaInstance } from '../types/ElysiaInstance.type';
+
+export const statusRoutes = (app: ElysiaInstance): ElysiaInstance => {
+  // =============================================================
+  //                     Get Job Status (/status/:jobId)
+  // =============================================================
+  app.get(
+    '/status/:jobId',
+    async ({ params }: { params: { jobId: string } }) => {
+      const { jobId } = params;
+      if (!jobId) {
+        return resp.BAD_REQUEST('Job ID is required.');
+      }
+      try {
+        const responsePayload = await getJobStatus(jobId);
+        return resp.SUCCESS(responsePayload);
+      } catch (error: any) {
+        console.error(`[API] Failed to get status for job ${jobId}:`, error);
+        return resp.ERROR('Failed to retrieve job status.' + error.message);
+      }
+    },
+    getStatusDoc
+  );
+
+  return app;
+};

packages/auth/src/index.ts

@@ -4,6 +4,11 @@ import * as authenticators from './lib/authenticators';
 import { WebAuthnAuthenticator } from './lib/authenticators';
 import { DiscordAuthenticator } from './lib/authenticators/native/DiscordAuthenticator';
 import { GoogleAuthenticator } from './lib/authenticators/native/GoogleAuthenticator';
+import { StytchEmailOtpAuthenticator } from './lib/authenticators/stytch/factors/StytchEmailOtpAuthenticator';
+import { StytchSmsOtpAuthenticator } from './lib/authenticators/stytch/factors/StytchSmsOtpAuthenticator';
+import { StytchTotp2FAAuthenticator } from './lib/authenticators/stytch/factors/2fa/StytchTotp2FAAuthenticator';
+
+import { StytchWhatsAppOtpAuthenticator } from './lib/authenticators/stytch/factors/StytchWhatsAppOtpAuthenticator';
 import { ViemAccountAuthenticator } from './lib/authenticators/ViemAccountAuthenticator';
 import { WalletClientAuthenticator } from './lib/authenticators/WalletClientAuthenticator';
 // import { GetAuthContext } from './lib/AuthManager/getAuthContext';
@@ -22,7 +27,7 @@ export type { LitAuthStorageProvider };
 /**
  * Type definition for the structure of authentication data used within the Lit Auth client.
  */
-  export type { LitAuthData };
+export type { LitAuthData };
 
 /**
  * Type definition for the structure of authentication context used within the Lit Auth client.
@@ -61,8 +66,8 @@ export { authenticators };
  */
 // export { createAuthManager } from './lib/auth-manager';
 // export { getAuthContext } from './lib/AuthManager/getAuthContext';
-  export { getEoaAuthContext } from './lib/AuthManager/authContexts/getEoaAuthContext';
-  export { getPkpAuthContext } from './lib/AuthManager/authContexts/getPkpAuthContext';
+export { getEoaAuthContext } from './lib/AuthManager/authContexts/getEoaAuthContext';
+export { getPkpAuthContext } from './lib/AuthManager/authContexts/getPkpAuthContext';
 /**
  * Class responsible for communicating with the Lit Relay server.
  * Used for operations like minting PKPs associated with authentication methods.
@@ -91,6 +96,11 @@ export {
   GoogleAuthenticator,
   ViemAccountAuthenticator,
   WalletClientAuthenticator,
-  WebAuthnAuthenticator
+  WebAuthnAuthenticator,
+  StytchEmailOtpAuthenticator,
+  StytchSmsOtpAuthenticator,
+  StytchTotp2FAAuthenticator,
+  StytchWhatsAppOtpAuthenticator,
 };
 
+export { ExampleAppAuthenticator } from './lib/authenticators/custom/ExampleAppAuthenticator';

packages/auth/src/lib/AuthManager/auth-manager.ts

@@ -1,8 +1,5 @@
 import { getChildLogger } from '@lit-protocol/logger';
-import {
-  AuthData,
-  HexPrefixedSchema
-} from '@lit-protocol/schemas';
+import { AuthData, HexPrefixedSchema } from '@lit-protocol/schemas';
 import { z } from 'zod';
 import { AuthConfigV2 } from '../authenticators/types';
 import type { LitAuthStorageProvider } from '../storage/types';
@@ -12,6 +9,9 @@ import {
 } from './authAdapters/getEoaAuthContextAdapter';
 import { getPkpAuthContextAdapter } from './authAdapters/getPkpAuthContextAdapter';
 import { AuthConfigSchema } from './authContexts/BaseAuthContextType';
+import { getCustomAuthContextAdapter } from './authAdapters/getCustomAuthContextAdapter';
+import { hexToBigInt, keccak256, toBytes } from 'viem';
+
 export interface AuthManagerParams {
   storage: LitAuthStorageProvider;
 }
@@ -76,14 +76,26 @@ export const createAuthManager = (authManagerParams: AuthManagerParams) => {
     }) => {
       return getPkpAuthContextAdapter(authManagerParams, params);
     },
-    // createCustomAuthContext: <T extends ICustomAuthenticator>(params: {
-    //   authenticator: T;
-    //   settings: ConstructorParameters<T>[0]; // Infer settings type from constructor
-    //   config: { pkpPublicKey: string; [key: string]: any }; // Execution config
-    //   authConfig: AuthConfigV2;
-    //   litClient: BaseAuthContext<any>['litClient'];
-    // }) => {
-    //   return getCustomAuthContextAdapter(authManagerParams, params);
-    // },
+    createCustomAuthContext: (params: {
+      // authData: AuthData;
+      pkpPublicKey: z.infer<typeof HexPrefixedSchema>;
+      authConfig: AuthConfigV2;
+      litClient: BaseAuthContext<any>['litClient'];
+
+      // custom auth params
+      customAuthParams: {
+        litActionCode?: string;
+        litActionIpfsId?: string;
+        jsParams?: Record<string, any>;
+      };
+    }) => {
+      // make jsParams nested inside jsParams so that
+      // the dev can check all variables inside an object in Lit action
+      params.customAuthParams.jsParams = {
+        jsParams: params.customAuthParams.jsParams,
+      };
+
+      return getCustomAuthContextAdapter(authManagerParams, params);
+    },
   };
 };

packages/auth/src/lib/AuthManager/authAdapters/getCustomAuthContextAdapter.ts

@@ -1,111 +1,107 @@
-import { ethers } from 'ethers';
+import { AUTH_METHOD_TYPE_VALUES, PRODUCT_IDS } from '@lit-protocol/constants';
 import {
-  AuthConfig,
-  AuthManagerParams,
-  BaseAuthContext,
-} from '../auth-manager';
-import { PkpAuthDepsSchema } from './getPkpAuthContextAdapter';
-import { AUTH_METHOD_TYPE } from '@lit-protocol/constants';
+  AuthData,
+  HexPrefixedSchema,
+  NodeUrlsSchema,
+} from '@lit-protocol/schemas';
+import { ethers } from 'ethers';
+import { z } from 'zod';
+import { AuthConfigV2 } from '../../authenticators/types';
+import { AuthManagerParams } from '../auth-manager';
+import { getCustomAuthContext } from '../authContexts/getCustomAuthContext';
+import { processResources } from '../utils/processResources';
 import { tryGetCachedAuthData } from '../try-getters/tryGetCachedAuthData';
 
-export interface ICustomAuthenticator {
-  new (settings: any): ICustomAuthenticatorInstance;
-  LIT_ACTION_CODE_BASE64?: string;
-  LIT_ACTION_IPFS_ID?: string;
-}
-
-interface ICustomAuthenticatorInstance {
-  // Method to perform external auth and return jsParams for the Lit Action
-  // Accepts the config object which includes pkpPublicKey and other needed params
-  authenticate(config: {
-    pkpPublicKey: string;
-    [key: string]: any;
-  }): Promise<Record<string, any> | null>;
-}
+export const CustomAuthDepsSchema = z.object({
+  nonce: z.any(),
+  currentEpoch: z.any(),
+  getSignSessionKey: z.any(),
+  nodeUrls: NodeUrlsSchema,
+});
 
 export async function getCustomAuthContextAdapter(
   upstreamParams: AuthManagerParams,
   params: {
-    authenticator: ICustomAuthenticator; // Use the interface type
-    settings: Record<string, any>; // For constructor
-    config: { pkpPublicKey: string; [key: string]: any }; // For authenticate method
-    authConfig: AuthConfig; // For SIWE/session
-    litClient: BaseAuthContext<any>['litClient'];
+    // authData: AuthData;
+    pkpPublicKey: z.infer<typeof HexPrefixedSchema>;
+    authConfig: AuthConfigV2;
+    litClient: {
+      getContext: () => Promise<any>;
+    };
+    customAuthParams: {
+      litActionCode?: string;
+      litActionIpfsId?: string;
+      jsParams?: Record<string, any>;
+    };
   }
 ) {
-  // 1. Instantiate the custom authenticator helper using 'settings'
-  const customAuthHelper = new params.authenticator(params.settings);
+  const _resources = processResources(params.authConfig.resources);
 
-  // 2. Call the helper's authenticate method using 'config'
-  if (!customAuthHelper.authenticate) {
-    throw new Error("Custom authenticator is missing 'authenticate' method.");
-  }
-  // Pass the entire config object to the authenticator's authenticate method
-  const jsParams = await customAuthHelper.authenticate(params.config);
-  if (!jsParams) {
-    throw new Error('Custom authenticator failed to produce jsParams.');
-  }
+  // TODO: 👇 The plan is to identify if the certain operations could be wrapped inside a single function
+  // where different network modules can provide their own implementations.
 
-  // 3. Get the static Lit Action code/ID from the authenticator class
-  const litActionCode = params.authenticator.LIT_ACTION_CODE_BASE64;
-  const litActionIpfsId = params.authenticator.LIT_ACTION_IPFS_ID; // Optional
-  if (!litActionCode && !litActionIpfsId) {
-    throw new Error(
-      'Custom authenticator is missing static LIT_ACTION_CODE_BASE64 or LIT_ACTION_IPFS_ID.'
-    );
-  }
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this!
+  const litClientCtx = await params.litClient.getContext();
+
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this! (This can be in both Naga and Datil)
+  const latestConnectionInfo = litClientCtx.latestConnectionInfo;
 
-  // 4. Extract pkpPublicKey (already available in params.config)
-  const pkpPublicKey = params.config.pkpPublicKey;
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this! (This can only be in Naga)
+  const nodePrices = latestConnectionInfo.priceFeedInfo.networkPrices;
 
-  // 5. Get node dependencies, session key etc.
-  const litClientConfig = PkpAuthDepsSchema.parse({
-    nonce: await params.litClient.getLatestBlockhash(),
-    currentEpoch: await params.litClient.getCurrentEpoch(),
-    getSignSessionKey: params.litClient.getSignSessionKey,
-    nodeUrls: await params.litClient.getMaxPricesForNodeProduct({
-      product: 'LIT_ACTION', // Or appropriate product
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this! (This can be in both Naga and Datil)
+  const handshakeResult = litClientCtx.handshakeResult;
+
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this! (This can be in both Naga and Datil)
+  const threshold = handshakeResult.threshold;
+
+  // TODO: ❗️THIS IS NOT TYPED - we have to fix this! (This can only be in Naga)
+  const nodeUrls = litClientCtx.getMaxPricesForNodeProduct({
+    nodePrices: nodePrices,
+    userMaxPrice: litClientCtx.getUserMaxPrice({
+      product: 'LIT_ACTION',
     }),
+    productId: PRODUCT_IDS['LIT_ACTION'],
+    numRequiredNodes: threshold,
   });
-  const pkpAddress = ethers.utils.computeAddress(pkpPublicKey);
-  const authData = await tryGetCachedAuthData({
+
+  const pkpAddress = ethers.utils.computeAddress(params.pkpPublicKey);
+
+  const litAuthData = await tryGetCachedAuthData({
     storage: upstreamParams.storage,
     address: pkpAddress,
-    expiration: params.authConfig.expiration,
-    type: AUTH_METHOD_TYPE.LitAction, // Session type remains LitAction
+    expiration: params.authConfig.expiration!,
+    type: 'custom' as unknown as AUTH_METHOD_TYPE_VALUES,
   });
 
-  // 6. Prepare the request body for the node signing function
-  const requestBodyForCustomAuth = {
-    sessionKey: authData.sessionKey.keyPair.publicKey,
-    pkpPublicKey: pkpPublicKey,
-    statement: params.authConfig.statement,
-    domain: params.authConfig.domain,
-    expiration: params.authConfig.expiration,
-    resources: params.authConfig.resources,
-    uri: authData.sessionKey.keyPair.publicKey,
-    nonce: litClientConfig.nonce,
-    ...(litActionCode && { code: litActionCode }),
-    ...(litActionIpfsId && { litActionIpfsId: litActionIpfsId }),
-    jsParams: jsParams, // Use the result from customAuthHelper.authenticate
-    authMethods: [],
-    epoch: litClientConfig.currentEpoch,
-    // ... other fields like curveType, signingScheme ...
-  };
-
-  // 7. Return the auth context object
-  return {
-    chain: 'ethereum',
-    pkpPublicKey: pkpPublicKey,
-    resources: params.authConfig.resources,
-    capabilityAuthSigs: params.authConfig.capabilityAuthSigs,
-    expiration: params.authConfig.expiration,
-    authNeededCallback: async () => {
-      const authSig = await litClientConfig.getSignSessionKey({
-        requestBody: requestBodyForCustomAuth,
-        nodeUrls: litClientConfig.nodeUrls.map((node: any) => node.url),
-      });
-      return authSig;
+  return getCustomAuthContext({
+    authentication: {
+      pkpPublicKey: params.pkpPublicKey,
+      // authData: {} as any,
     },
-  };
+    authConfig: {
+      domain: params.authConfig.domain!,
+      resources: _resources,
+      capabilityAuthSigs: params.authConfig.capabilityAuthSigs!,
+      expiration: params.authConfig.expiration!,
+      statement: params.authConfig.statement!,
+    },
+    customParams: {
+      litActionCode: params.customAuthParams.litActionCode,
+      litActionIpfsId: params.customAuthParams.litActionIpfsId,
+      jsParams: params.customAuthParams.jsParams,
+    },
+    deps: {
+      litAuthData: litAuthData,
+      connection: {
+        nonce: litClientCtx.latestBlockhash,
+        currentEpoch:
+          litClientCtx.latestConnectionInfo.epochState.currentNumber,
+        nodeUrls: nodeUrls,
+      },
+      signCustomSessionKey: litClientCtx.signCustomSessionKey,
+      storage: upstreamParams.storage,
+      pkpAddress: pkpAddress,
+    },
+  });
 }