feat: remove keys on mediator when deleting connections (#1143)

Signed-off-by: Ariel Gentile <[email protected]>

packages/core/src/modules/connections/ConnectionsApi.ts

@@ -321,6 +321,19 @@ export class ConnectionsApi {
    * @param connectionId the connection record id
    */
   public async deleteById(connectionId: string) {
+    const connection = await this.connectionService.getById(this.agentContext, connectionId)
+
+    if (connection.mediatorId && connection.did) {
+      const did = await this.didResolverService.resolve(this.agentContext, connection.did)
+
+      if (did.didDocument) {
+        await this.routingService.removeRouting(this.agentContext, {
+          recipientKeys: did.didDocument.recipientKeys,
+          mediatorId: connection.mediatorId,
+        })
+      }
+    }
+
     return this.connectionService.deleteById(this.agentContext, connectionId)
   }
 

packages/core/src/modules/oob/OutOfBandApi.ts

@@ -1,6 +1,5 @@
 import type { AgentMessage } from '../../agent/AgentMessage'
 import type { AgentMessageReceivedEvent } from '../../agent/Events'
-import type { Key } from '../../crypto'
 import type { Attachment } from '../../decorators/attachment/Attachment'
 import type { Query } from '../../storage/StorageService'
 import type { PlaintextMessage } from '../../types'
@@ -16,6 +15,7 @@ import { filterContextCorrelationId, AgentEventTypes } from '../../agent/Events'
 import { MessageSender } from '../../agent/MessageSender'
 import { OutboundMessageContext } from '../../agent/models'
 import { InjectionSymbols } from '../../constants'
+import { Key } from '../../crypto'
 import { ServiceDecorator } from '../../decorators/service/ServiceDecorator'
 import { AriesFrameworkError } from '../../error'
 import { Logger } from '../../logger'
@@ -416,7 +416,6 @@ export class OutOfBandApi {
       label?: string
       alias?: string
       imageUrl?: string
-      mediatorId?: string
       routing?: Routing
     }
   ) {
@@ -565,6 +564,22 @@ export class OutOfBandApi {
    * @param outOfBandId the out of band record id
    */
   public async deleteById(outOfBandId: string) {
+    const outOfBandRecord = await this.getById(outOfBandId)
+
+    const relatedConnections = await this.connectionsApi.findAllByOutOfBandId(outOfBandId)
+
+    // If it uses mediation and there are no related connections, proceed to delete keys from mediator
+    // Note: if OOB Record is reusable, it is safe to delete it because every connection created from
+    // it will use its own recipient key
+    if (outOfBandRecord.mediatorId && (relatedConnections.length === 0 || outOfBandRecord.reusable)) {
+      const recipientKeys = outOfBandRecord.getTags().recipientKeyFingerprints.map((item) => Key.fromFingerprint(item))
+
+      await this.routingService.removeRouting(this.agentContext, {
+        recipientKeys,
+        mediatorId: outOfBandRecord.mediatorId,
+      })
+    }
+
     return this.outOfBandService.deleteById(this.agentContext, outOfBandId)
   }
 

packages/core/src/modules/routing/RecipientApi.ts

@@ -18,8 +18,10 @@ import { AriesFrameworkError } from '../../error'
 import { Logger } from '../../logger'
 import { inject, injectable } from '../../plugins'
 import { TransportEventTypes } from '../../transport'
+import { ConnectionMetadataKeys } from '../connections/repository/ConnectionMetadataTypes'
 import { ConnectionService } from '../connections/services'
 import { DidsApi } from '../dids'
+import { verkeyToDidKey } from '../dids/helpers'
 import { DiscoverFeaturesApi } from '../discover-features'
 
 import { MediatorPickupStrategy } from './MediatorPickupStrategy'
@@ -28,6 +30,7 @@ import { RoutingEventTypes } from './RoutingEvents'
 import { KeylistUpdateResponseHandler } from './handlers/KeylistUpdateResponseHandler'
 import { MediationDenyHandler } from './handlers/MediationDenyHandler'
 import { MediationGrantHandler } from './handlers/MediationGrantHandler'
+import { KeylistUpdate, KeylistUpdateAction, KeylistUpdateMessage } from './messages'
 import { MediationState } from './models/MediationState'
 import { StatusRequestMessage, BatchPickupMessage, StatusMessage } from './protocol'
 import { StatusHandler, MessageDeliveryHandler } from './protocol/pickup/v2/handlers'
@@ -370,8 +373,22 @@ export class RecipientApi {
     return mediationRecord
   }
 
-  public async notifyKeylistUpdate(connection: ConnectionRecord, verkey: string) {
-    const message = this.mediationRecipientService.createKeylistUpdateMessage(verkey)
+  public async notifyKeylistUpdate(connection: ConnectionRecord, verkey: string, action?: KeylistUpdateAction) {
+    // Use our useDidKey configuration unless we know the key formatting other party is using
+    let useDidKey = this.agentContext.config.useDidKeyInProtocols
+
+    const useDidKeysConnectionMetadata = connection.metadata.get(ConnectionMetadataKeys.UseDidKeysForProtocol)
+    if (useDidKeysConnectionMetadata) {
+      useDidKey = useDidKeysConnectionMetadata[KeylistUpdateMessage.type.protocolUri] ?? useDidKey
+    }
+
+    const message = this.mediationRecipientService.createKeylistUpdateMessage([
+      new KeylistUpdate({
+        action: action ?? KeylistUpdateAction.add,
+        recipientKey: useDidKey ? verkeyToDidKey(verkey) : verkey,
+      }),
+    ])
+
     const outboundMessageContext = new OutboundMessageContext(message, {
       agentContext: this.agentContext,
       connection,

packages/core/src/modules/routing/messages/KeylistUpdateMessage.ts

@@ -9,8 +9,13 @@ export enum KeylistUpdateAction {
   remove = 'remove',
 }
 
+export interface KeylistUpdateOptions {
+  recipientKey: string
+  action: KeylistUpdateAction
+}
+
 export class KeylistUpdate {
-  public constructor(options: { recipientKey: string; action: KeylistUpdateAction }) {
+  public constructor(options: KeylistUpdateOptions) {
     if (options) {
       this.recipientKey = options.recipientKey
       this.action = options.action

packages/core/src/modules/routing/services/MediationRecipientService.ts

@@ -9,7 +9,7 @@ import type { Routing } from '../../connections/services/ConnectionService'
 import type { MediationStateChangedEvent, KeylistUpdatedEvent } from '../RoutingEvents'
 import type { MediationDenyMessage } from '../messages'
 import type { StatusMessage, MessageDeliveryMessage } from '../protocol'
-import type { GetRoutingOptions } from './RoutingService'
+import type { GetRoutingOptions, RemoveRoutingOptions } from './RoutingService'
 
 import { firstValueFrom, ReplaySubject } from 'rxjs'
 import { filter, first, timeout } from 'rxjs/operators'
@@ -25,7 +25,8 @@ import { JsonTransformer } from '../../../utils'
 import { ConnectionType } from '../../connections/models/ConnectionType'
 import { ConnectionMetadataKeys } from '../../connections/repository/ConnectionMetadataTypes'
 import { ConnectionService } from '../../connections/services/ConnectionService'
-import { didKeyToVerkey, isDidKey, verkeyToDidKey } from '../../dids/helpers'
+import { DidKey } from '../../dids'
+import { didKeyToVerkey, isDidKey } from '../../dids/helpers'
 import { ProblemReportError } from '../../problem-reports'
 import { RecipientModuleConfig } from '../RecipientModuleConfig'
 import { RoutingEventTypes } from '../RoutingEvents'
@@ -174,7 +175,7 @@ export class MediationRecipientService {
   public async keylistUpdateAndAwait(
     agentContext: AgentContext,
     mediationRecord: MediationRecord,
-    verKey: string,
+    updates: { recipientKey: Key; action: KeylistUpdateAction }[],
     timeoutMs = 15000 // TODO: this should be a configurable value in agent config
   ): Promise<MediationRecord> {
     const connection = await this.connectionService.getById(agentContext, mediationRecord.connectionId)
@@ -187,7 +188,15 @@ export class MediationRecipientService {
       useDidKey = useDidKeysConnectionMetadata[KeylistUpdateMessage.type.protocolUri] ?? useDidKey
     }
 
-    const message = this.createKeylistUpdateMessage(useDidKey ? verkeyToDidKey(verKey) : verKey)
+    const message = this.createKeylistUpdateMessage(
+      updates.map(
+        (item) =>
+          new KeylistUpdate({
+            action: item.action,
+            recipientKey: useDidKey ? new DidKey(item.recipientKey).did : item.recipientKey.publicKeyBase58,
+          })
+      )
+    )
 
     mediationRecord.assertReady()
     mediationRecord.assertRole(MediationRole.Recipient)
@@ -216,14 +225,9 @@ export class MediationRecipientService {
     return keylistUpdate.payload.mediationRecord
   }
 
-  public createKeylistUpdateMessage(verkey: string): KeylistUpdateMessage {
+  public createKeylistUpdateMessage(updates: KeylistUpdate[]): KeylistUpdateMessage {
     const keylistUpdateMessage = new KeylistUpdateMessage({
-      updates: [
-        new KeylistUpdate({
-          action: KeylistUpdateAction.add,
-          recipientKey: verkey,
-        }),
-      ],
+      updates,
     })
     return keylistUpdateMessage
   }
@@ -247,19 +251,43 @@ export class MediationRecipientService {
     if (!mediationRecord) return routing
 
     // new did has been created and mediator needs to be updated with the public key.
-    mediationRecord = await this.keylistUpdateAndAwait(
-      agentContext,
-      mediationRecord,
-      routing.recipientKey.publicKeyBase58
-    )
+    mediationRecord = await this.keylistUpdateAndAwait(agentContext, mediationRecord, [
+      {
+        recipientKey: routing.recipientKey,
+        action: KeylistUpdateAction.add,
+      },
+    ])
 
     return {
       ...routing,
+      mediatorId: mediationRecord.id,
       endpoints: mediationRecord.endpoint ? [mediationRecord.endpoint] : routing.endpoints,
       routingKeys: mediationRecord.routingKeys.map((key) => Key.fromPublicKeyBase58(key, KeyType.Ed25519)),
     }
   }
 
+  public async removeMediationRouting(
+    agentContext: AgentContext,
+    { recipientKeys, mediatorId }: RemoveRoutingOptions
+  ): Promise<void> {
+    const mediationRecord = await this.getById(agentContext, mediatorId)
+
+    if (!mediationRecord) {
+      throw new AriesFrameworkError('No mediation record to remove routing from has been found')
+    }
+
+    await this.keylistUpdateAndAwait(
+      agentContext,
+      mediationRecord,
+      recipientKeys.map((item) => {
+        return {
+          recipientKey: item,
+          action: KeylistUpdateAction.remove,
+        }
+      })
+    )
+  }
+
   public async processMediationDeny(messageContext: InboundMessageContext<MediationDenyMessage>) {
     const connection = messageContext.assertReadyConnection()
 

packages/core/src/modules/routing/services/RoutingService.ts

@@ -52,6 +52,10 @@ export class RoutingService {
 
     return routing
   }
+
+  public async removeRouting(agentContext: AgentContext, options: RemoveRoutingOptions) {
+    await this.mediationRecipientService.removeMediationRouting(agentContext, options)
+  }
 }
 
 export interface GetRoutingOptions {
@@ -66,3 +70,15 @@ export interface GetRoutingOptions {
    */
   useDefaultMediator?: boolean
 }
+
+export interface RemoveRoutingOptions {
+  /**
+   * Keys to remove routing from
+   */
+  recipientKeys: Key[]
+
+  /**
+   * Identifier of the mediator used when routing has been set up
+   */
+  mediatorId: string
+}