feat: drain tracked NWC handlers on stop

src/lnurlcash/nwc/service.ts

@@ -14,26 +14,18 @@
 // foreground-only design itself is documented in the nwc.ts façade
 // header.
 
+import {linkingPubKeyHex} from '../keys'
 import type {NwcConnectionRecord} from '../storage/nwcConnections'
 import {readNwcConnections} from '../storage/nwcConnections'
+import {assertSavedKeyOwner} from '../storage/currentOwner'
 
 import type {NwcConnectionInfo} from './connection'
 import {deriveNwcWalletKey, nwcWalletPubkey} from './connection'
 import type {NwcServiceDeps, PendingInvoice, RequestContext} from './context'
 import {dispatch} from './dispatch'
-import type {
-  NostrEvent,
-  NwcEncryption,
-  NwcRequest,
-  NwcResponse
-} from './protocol'
-import {
-  NWC_REQUEST_KIND,
-  buildInfoEvent,
-  buildResponseEvent,
-  decryptRequest
-} from './protocol'
-import type {NwcSubscription, NwcTransport} from './transport'
+import type {NostrEvent, NwcEncryption, NwcRequest, NwcResponse} from './protocol'
+import {NWC_REQUEST_KIND, buildInfoEvent, buildResponseEvent, decryptRequest} from './protocol'
+import type {NwcSubscription} from './transport'
 import {defaultNwcTransport} from './transport'
 
 export type {NwcConnectionInfo}
@@ -44,7 +36,8 @@ const MAX_REQUEST_AGE_SECONDS = 600
 
 type ConnectionRuntime = {
   info: NwcConnectionInfo
-  walletSecret: Uint8Array
+  // nulled and zeroed only after every tracked handler drains
+  walletSecret: Uint8Array | null
   // invoices this connection issued, by payment hash - in-memory only:
   // pending invoices don't survive a restart (lookup then answers
   // NOT_FOUND), same as any foreground-only wallet
@@ -61,34 +54,38 @@ export type NwcService = {
   connections: NwcConnectionInfo[]
   // closes every relay subscription. In-flight handlers still finish -
   // their changesets hold money - but no new requests are picked up
-  stop: () => void
+  stop: () => Promise<void>
 }
 
 export const startService = async (
   linkingPrivKey: Uint8Array,
   deps: NwcServiceDeps,
-  records: NwcConnectionRecord[] = readNwcConnections()
+  records?: NwcConnectionRecord[],
 ): Promise<NwcService> => {
   const transport = deps.transport ?? (await defaultNwcTransport())
-  const nowSeconds = (): number =>
-    deps.nowSeconds?.() ?? Math.floor(Date.now() / 1000)
+  const ownerId = linkingPubKeyHex(linkingPrivKey)
+  const ownedRecords = (records ?? readNwcConnections(ownerId)).filter(
+    (record) => record.ownerId === ownerId,
+  )
+  const nowSeconds = (): number => deps.nowSeconds?.() ?? Math.floor(Date.now() / 1000)
 
   const publishResponse = async (
     runtime: ConnectionRuntime,
+    walletSecret: Uint8Array,
     requestEventId: string,
     encryption: NwcEncryption,
-    response: NwcResponse
+    response: NwcResponse,
   ): Promise<void> => {
     await transport.publish(
       runtime.info.record.relays,
       buildResponseEvent(
-        runtime.walletSecret,
+        walletSecret,
         runtime.info.record.clientPubkey,
         encryption,
         requestEventId,
         response,
-        nowSeconds()
-      )
+        nowSeconds(),
+      ),
     )
   }
 
@@ -97,12 +94,12 @@ export const startService = async (
   const dispatchSerialized = (
     runtime: ConnectionRuntime,
     ctx: RequestContext,
-    request: NwcRequest
+    request: NwcRequest,
   ): Promise<NwcResponse> => {
     if (request.method !== 'pay_invoice') return dispatch(ctx, request)
     const run = runtime.queue.then(
       () => dispatch(ctx, request),
-      () => dispatch(ctx, request)
+      () => dispatch(ctx, request),
     )
     runtime.queue = run.catch(() => undefined)
     return run
@@ -111,28 +108,53 @@ export const startService = async (
   const handleEvent = async (
     runtime: ConnectionRuntime,
     ctx: RequestContext,
-    event: NostrEvent
+    event: NostrEvent,
   ): Promise<void> => {
+    const walletSecret = runtime.walletSecret
+    if (walletSecret === null) return
     const at = nowSeconds()
     // replay safety (see the header): too-old requests are dropped
     if (event.created_at < at - MAX_REQUEST_AGE_SECONDS) return
     const decrypted = decryptRequest(
-      runtime.walletSecret,
+      walletSecret,
       runtime.info.walletServicePubkey,
       runtime.info.record.clientPubkey,
       event,
-      at
+      at,
     )
     if (decrypted === null) return
     if (decrypted.respond) {
-      await publishResponse(runtime, event.id, decrypted.encryption, decrypted.response)
+      await publishResponse(
+        runtime,
+        walletSecret,
+        event.id,
+        decrypted.encryption,
+        decrypted.response,
+      )
       return
     }
     const response = await dispatchSerialized(runtime, ctx, decrypted.request)
-    await publishResponse(runtime, event.id, decrypted.encryption, response)
+    await publishResponse(runtime, walletSecret, event.id, decrypted.encryption, response)
   }
 
-  const runtimes = records.map(record => {
+  let accepting = true
+  // interrupts long observation waits (the invoice claim poll) at stop;
+  // the drain below still awaits tasks that reached a fund-critical commit
+  const stopController = new AbortController()
+  const inFlight = new Set<Promise<void>>()
+  const track = (task: Promise<void>): void => {
+    inFlight.add(task)
+    void task.then(
+      () => inFlight.delete(task),
+      () => inFlight.delete(task),
+    )
+  }
+  const startBackground = (work: () => Promise<void>): boolean => {
+    if (!accepting) return false
+    track(work())
+    return true
+  }
+  let runtimes = ownedRecords.map((record) => {
     const walletSecret = deriveNwcWalletKey(linkingPrivKey, record.clientPubkey)
     const runtime: ConnectionRuntime = {
       info: {record, walletServicePubkey: nwcWalletPubkey(walletSecret)},
@@ -141,29 +163,37 @@ export const startService = async (
       queue: Promise.resolve(),
       // replaced below, immediately - the field exists because the
       // subscription callback closes over the runtime
-      sub: {close: () => undefined}
+      sub: {close: () => undefined},
     }
     const ctx: RequestContext = {
       deps,
       connection: () => runtime.info,
-      updateRecord: updated => {
+      updateRecord: (updated) => {
         runtime.info = {...runtime.info, record: updated}
       },
       invoices: runtime.invoices,
-      nowSeconds
+      nowSeconds,
+      assertOwner: () => {
+        deps.assertCurrentOwner()
+        assertSavedKeyOwner(ownerId)
+      },
+      startBackground,
+      stopSignal: stopController.signal,
     }
     runtime.sub = transport.subscribe(
       record.relays,
       {
         kinds: [NWC_REQUEST_KIND],
         '#p': [runtime.info.walletServicePubkey],
-        since: nowSeconds()
+        since: nowSeconds(),
       },
-      event => {
-        void handleEvent(runtime, ctx, event).catch(err =>
+      (event) => {
+        if (!accepting) return
+        const handler = handleEvent(runtime, ctx, event).catch((err) => {
           deps.onError?.(err, runtime.info)
-        )
-      }
+        })
+        track(handler)
+      },
     )
     return runtime
   })
@@ -171,20 +201,43 @@ export const startService = async (
   // info events: best-effort - a rejected publish must not sink startup;
   // the client learns capabilities from its first error-free exchange too
   for (const runtime of runtimes) {
+    const walletSecret = runtime.walletSecret
+    if (walletSecret === null) continue
     try {
       await transport.publish(
         runtime.info.record.relays,
-        buildInfoEvent(runtime.walletSecret, nowSeconds())
+        buildInfoEvent(walletSecret, nowSeconds()),
+      )
+    } catch (error) {
+      deps.onError?.(
+        error instanceof Error ? error : new Error('NWC info publication failed.', {cause: error}),
+        runtime.info,
       )
-    } catch (err) {
-      deps.onError?.(err, runtime.info)
     }
   }
 
-  return {
-    connections: runtimes.map(r => r.info),
-    stop: () => {
+  const connections = runtimes.map((runtime) => runtime.info)
+  let stopPromise: Promise<void> | null = null
+  const stop = (): Promise<void> => {
+    if (stopPromise !== null) return stopPromise
+    accepting = false
     for (const runtime of runtimes) runtime.sub.close()
+    stopController.abort()
+    stopPromise = Promise.all([...inFlight])
+      .then(() => undefined)
+      .finally(() => {
+        for (const runtime of runtimes) {
+          runtime.walletSecret?.fill(0)
+          runtime.walletSecret = null
+          runtime.invoices.clear()
         }
+        runtimes = []
+      })
+    return stopPromise
+  }
+
+  return {
+    connections,
+    stop,
   }
 }

src/lnurlcash/nwc/transport.ts

@@ -19,7 +19,7 @@ export type NwcTransport = {
   subscribe: (
     relays: string[],
     filter: NostrFilter,
-    onEvent: (event: NostrEvent) => void
+    onEvent: (event: NostrEvent) => void,
   ) => NwcSubscription
 }
 
@@ -30,11 +30,10 @@ export const defaultNwcTransport = async (): Promise<NwcTransport> => {
     publish: async (relays, event) => {
       const results = await Promise.allSettled(pool.publish(relays, event))
       // one honest relay accepting is enough - same rule as the backup
-      if (!results.some(r => r.status === 'fulfilled')) {
+      if (!results.some((r) => r.status === 'fulfilled')) {
         throw new Error('No relay accepted the event.')
       }
     },
-    subscribe: (relays, filter, onEvent) =>
-      pool.subscribeMany(relays, filter, {onevent: onEvent})
+    subscribe: (relays, filter, onEvent) => pool.subscribeMany(relays, filter, {onevent: onEvent}),
   }
 }