Pass request to PromptCodeHandler in guard registerUI (#909)

* Pass request to PromptCodeHandler in guard registerUI

* Fixing guard registerUI test

Author: corbanbrook

packages/wallet/wdk/src/sequence/handlers/guard.ts

@@ -5,21 +5,25 @@ import { BaseSignatureRequest, SignerUnavailable, SignerReady, SignerActionable,
 import { Signatures } from '../signatures.js'
 import { Guards } from '../guards.js'
 
+type RespondFn = (code: string) => Promise<void>
+
+export type PromptCodeHandler = (
+  request: BaseSignatureRequest,
+  codeType: 'TOTP' | 'PIN',
+  respond: RespondFn,
+) => Promise<void>
+
 export class GuardHandler implements Handler {
   kind = Kinds.Guard
 
-  private onPromptCode:
-    | undefined
-    | ((codeType: 'TOTP' | 'PIN', respond: (code: string) => Promise<void>) => Promise<void>)
+  private onPromptCode: undefined | PromptCodeHandler
 
   constructor(
     private readonly signatures: Signatures,
     private readonly guards: Guards,
   ) {}
 
-  public registerUI(
-    onPromptCode: (codeType: 'TOTP' | 'PIN', respond: (code: string) => Promise<void>) => Promise<void>,
-  ) {
+  public registerUI(onPromptCode: PromptCodeHandler) {
     this.onPromptCode = onPromptCode
     return () => {
       this.onPromptCode = undefined
@@ -90,7 +94,7 @@ export class GuardHandler implements Handler {
             resolve(true)
           } catch (e) {
             if (e instanceof Guard.AuthRequiredError) {
-              const respond = async (code: string) => {
+              const respond: RespondFn = async (code: string) => {
                 try {
                   const signature = await guard.signEnvelope(request.envelope, { id: e.id, code })
                   await this.signatures.addSignature(request.id, signature)
@@ -100,7 +104,7 @@ export class GuardHandler implements Handler {
                 }
               }
 
-              await onPromptCode(e.id, respond)
+              await onPromptCode(request, e.id, respond)
             } else {
               reject(e)
             }

packages/wallet/wdk/src/sequence/handlers/mnemonic.ts

@@ -7,15 +7,17 @@ import { SignerReady, SignerUnavailable, BaseSignatureRequest, SignerActionable
 
 type RespondFn = (mnemonic: string) => Promise<void>
 
+export type PromptMnemonicHandler = (respond: RespondFn) => Promise<void>
+
 export class MnemonicHandler implements Handler {
   kind = Kinds.LoginMnemonic
 
-  private onPromptMnemonic: undefined | ((respond: RespondFn) => Promise<void>)
+  private onPromptMnemonic: undefined | PromptMnemonicHandler
   private readySigners = new Map<Address.Address, Signers.Pk.Pk>()
 
   constructor(private readonly signatures: Signatures) {}
 
-  public registerUI(onPromptMnemonic: (respond: RespondFn) => Promise<void>) {
+  public registerUI(onPromptMnemonic: PromptMnemonicHandler) {
     this.onPromptMnemonic = onPromptMnemonic
     return () => {
       this.onPromptMnemonic = undefined
@@ -93,7 +95,7 @@ export class MnemonicHandler implements Handler {
       message: 'enter-mnemonic',
       handle: () =>
         new Promise(async (resolve, reject) => {
-          const respond = async (mnemonic: string) => {
+          const respond: RespondFn = async (mnemonic: string) => {
             const signer = MnemonicHandler.toSigner(mnemonic)
             if (!signer) {
               return reject('invalid-mnemonic')

packages/wallet/wdk/src/sequence/handlers/otp.ts

@@ -11,16 +11,18 @@ import { AnswerIncorrectError, ChallengeExpiredError, TooManyAttemptsError } fro
 
 type RespondFn = (otp: string) => Promise<void>
 
+export type PromptOtpHandler = (recipient: string, respond: RespondFn) => Promise<void>
+
 export class OtpHandler extends IdentityHandler implements Handler {
   kind = Kinds.LoginEmailOtp
 
-  private onPromptOtp: undefined | ((recipient: string, respond: RespondFn) => Promise<void>)
+  private onPromptOtp: undefined | PromptOtpHandler
 
   constructor(nitro: Identity.IdentityInstrument, signatures: Signatures, authKeys: Db.AuthKeys) {
     super(nitro, authKeys, signatures, Identity.IdentityType.Email)
   }
 
-  public registerUI(onPromptOtp: (recipient: string, respond: RespondFn) => Promise<void>) {
+  public registerUI(onPromptOtp: PromptOtpHandler) {
     this.onPromptOtp = onPromptOtp
     return () => {
       this.onPromptOtp = undefined
@@ -91,13 +93,13 @@ export class OtpHandler extends IdentityHandler implements Handler {
 
   private handleAuth(
     challenge: Identity.OtpChallenge,
-    onPromptOtp: (recipient: string, respond: RespondFn) => Promise<void>,
+    onPromptOtp: PromptOtpHandler,
   ): Promise<{ signer: Signers.Signer & Signers.Witnessable; email: string }> {
     return new Promise(async (resolve, reject) => {
       try {
         const { loginHint, challenge: codeChallenge } = await this.nitroCommitVerifier(challenge)
 
-        const respond = async (otp: string) => {
+        const respond: RespondFn = async (otp: string) => {
           try {
             const { signer, email: returnedEmail } = await this.nitroCompleteAuth(
               challenge.withAnswer(codeChallenge, otp),

packages/wallet/wdk/src/sequence/manager.ts

@@ -27,8 +27,10 @@ import { Signers } from './signers.js'
 import { Transactions, TransactionsInterface } from './transactions.js'
 import { Kinds } from './types/signer.js'
 import { Wallets, WalletsInterface } from './wallets.js'
-import { GuardHandler } from './handlers/guard.js'
+import { GuardHandler, PromptCodeHandler } from './handlers/guard.js'
 import { PasskeyCredential } from '../dbs/index.js'
+import { PromptMnemonicHandler } from './handlers/mnemonic.js'
+import { PromptOtpHandler } from './handlers/otp.js'
 
 export type ManagerOptions = {
   verbose?: boolean
@@ -500,18 +502,16 @@ export class Manager {
     }
   }
 
-  public registerMnemonicUI(onPromptMnemonic: (respond: (mnemonic: string) => Promise<void>) => Promise<void>) {
+  public registerMnemonicUI(onPromptMnemonic: PromptMnemonicHandler) {
     return this.mnemonicHandler.registerUI(onPromptMnemonic)
   }
 
-  public registerOtpUI(onPromptOtp: (recipient: string, respond: (otp: string) => Promise<void>) => Promise<void>) {
+  public registerOtpUI(onPromptOtp: PromptOtpHandler) {
     return this.otpHandler?.registerUI(onPromptOtp) || (() => {})
   }
 
-  public registerGuardUI(
-    onPromptOtp: (codeType: 'TOTP' | 'PIN', respond: (otp: string) => Promise<void>) => Promise<void>,
-  ) {
-    return this.guardHandler?.registerUI(onPromptOtp) || (() => {})
+  public registerGuardUI(onPromptCode: PromptCodeHandler) {
+    return this.guardHandler?.registerUI(onPromptCode) || (() => {})
   }
 
   public async setRedirectPrefix(prefix: string) {

packages/wallet/wdk/test/guard.test.ts

@@ -223,7 +223,7 @@ describe('GuardHandler', () => {
       const mockAddSignature = vi.fn()
       signatures.addSignature = mockAddSignature
 
-      const mockCallback = vi.fn().mockImplementation(async (codeType, respond) => {
+      const mockCallback = vi.fn().mockImplementation(async (request, codeType, respond) => {
         expect(codeType).toBe('TOTP')
         await respond('123456')
       })