fix: connection id in sessions for new connections (openwallet-founda…

…tion#1383)

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

packages/core/src/agent/TransportService.ts

@@ -4,6 +4,7 @@ import type { DidDocument } from '../modules/dids'
 import type { EncryptedMessage } from '../types'
 
 import { DID_COMM_TRANSPORT_QUEUE } from '../constants'
+import { AriesFrameworkError } from '../error'
 import { injectable } from '../plugins'
 
 @injectable()
@@ -18,6 +19,15 @@ export class TransportService {
     return Object.values(this.transportSessionTable).find((session) => session?.connectionId === connectionId)
   }
 
+  public setConnectionIdForSession(sessionId: string, connectionId: string) {
+    const session = this.findSessionById(sessionId)
+    if (!session) {
+      throw new AriesFrameworkError(`Session not found with id ${sessionId}`)
+    }
+    session.connectionId = connectionId
+    this.saveSession(session)
+  }
+
   public hasInboundEndpoint(didDocument: DidDocument): boolean {
     return Boolean(didDocument.service?.find((s) => s.serviceEndpoint !== DID_COMM_TRANSPORT_QUEUE))
   }

packages/core/src/modules/connections/handlers/ConnectionRequestHandler.ts

@@ -5,6 +5,7 @@ import type { RoutingService } from '../../routing/services/RoutingService'
 import type { ConnectionsModuleConfig } from '../ConnectionsModuleConfig'
 import type { ConnectionService } from '../services/ConnectionService'
 
+import { TransportService } from '../../../agent/TransportService'
 import { OutboundMessageContext } from '../../../agent/models'
 import { AriesFrameworkError } from '../../../error/AriesFrameworkError'
 import { tryParseDid } from '../../dids/domain/parse'
@@ -34,7 +35,7 @@ export class ConnectionRequestHandler implements MessageHandler {
   }
 
   public async handle(messageContext: MessageHandlerInboundMessage<ConnectionRequestHandler>) {
-    const { agentContext, connection, recipientKey, senderKey, message } = messageContext
+    const { agentContext, connection, recipientKey, senderKey, message, sessionId } = messageContext
 
     if (!recipientKey || !senderKey) {
       throw new AriesFrameworkError('Unable to process connection request without senderVerkey or recipientKey')
@@ -62,30 +63,31 @@ export class ConnectionRequestHandler implements MessageHandler {
       )
     }
 
-    const receivedDidRecord = await this.didRepository.findReceivedDidByRecipientKey(
-      messageContext.agentContext,
-      senderKey
-    )
+    const receivedDidRecord = await this.didRepository.findReceivedDidByRecipientKey(agentContext, senderKey)
     if (receivedDidRecord) {
       throw new AriesFrameworkError(`A received did record for sender key ${senderKey.fingerprint} already exists.`)
     }
 
     const connectionRecord = await this.connectionService.processRequest(messageContext, outOfBandRecord)
 
+    // Associate the new connection with the session created for the inbound message
+    if (sessionId) {
+      const transportService = agentContext.dependencyManager.resolve(TransportService)
+      transportService.setConnectionIdForSession(sessionId, connectionRecord.id)
+    }
+
     if (connectionRecord?.autoAcceptConnection ?? this.connectionsModuleConfig.autoAcceptConnections) {
       // TODO: Allow rotation of keys used in the invitation for new ones not only when out-of-band is reusable
-      const routing = outOfBandRecord.reusable
-        ? await this.routingService.getRouting(messageContext.agentContext)
-        : undefined
+      const routing = outOfBandRecord.reusable ? await this.routingService.getRouting(agentContext) : undefined
 
       const { message } = await this.connectionService.createResponse(
-        messageContext.agentContext,
+        agentContext,
         connectionRecord,
         outOfBandRecord,
         routing
       )
       return new OutboundMessageContext(message, {
-        agentContext: messageContext.agentContext,
+        agentContext,
         connection: connectionRecord,
         outOfBand: outOfBandRecord,
       })

packages/core/src/modules/connections/handlers/DidExchangeRequestHandler.ts

@@ -5,6 +5,7 @@ import type { RoutingService } from '../../routing/services/RoutingService'
 import type { ConnectionsModuleConfig } from '../ConnectionsModuleConfig'
 import type { DidExchangeProtocol } from '../DidExchangeProtocol'
 
+import { TransportService } from '../../../agent/TransportService'
 import { OutboundMessageContext } from '../../../agent/models'
 import { AriesFrameworkError } from '../../../error/AriesFrameworkError'
 import { tryParseDid } from '../../dids/domain/parse'
@@ -35,7 +36,7 @@ export class DidExchangeRequestHandler implements MessageHandler {
   }
 
   public async handle(messageContext: MessageHandlerInboundMessage<DidExchangeRequestHandler>) {
-    const { agentContext, recipientKey, senderKey, message, connection } = messageContext
+    const { agentContext, recipientKey, senderKey, message, connection, sessionId } = messageContext
 
     if (!recipientKey || !senderKey) {
       throw new AriesFrameworkError('Unable to process connection request without senderKey or recipientKey')
@@ -65,10 +66,7 @@ export class DidExchangeRequestHandler implements MessageHandler {
       )
     }
 
-    const receivedDidRecord = await this.didRepository.findReceivedDidByRecipientKey(
-      messageContext.agentContext,
-      senderKey
-    )
+    const receivedDidRecord = await this.didRepository.findReceivedDidByRecipientKey(agentContext, senderKey)
     if (receivedDidRecord) {
       throw new AriesFrameworkError(`A received did record for sender key ${senderKey.fingerprint} already exists.`)
     }
@@ -83,21 +81,25 @@ export class DidExchangeRequestHandler implements MessageHandler {
 
     const connectionRecord = await this.didExchangeProtocol.processRequest(messageContext, outOfBandRecord)
 
+    // Associate the new connection with the session created for the inbound message
+    if (sessionId) {
+      const transportService = agentContext.dependencyManager.resolve(TransportService)
+      transportService.setConnectionIdForSession(sessionId, connectionRecord.id)
+    }
+
     if (connectionRecord.autoAcceptConnection ?? this.connectionsModuleConfig.autoAcceptConnections) {
       // TODO We should add an option to not pass routing and therefore do not rotate keys and use the keys from the invitation
       // TODO: Allow rotation of keys used in the invitation for new ones not only when out-of-band is reusable
-      const routing = outOfBandRecord.reusable
-        ? await this.routingService.getRouting(messageContext.agentContext)
-        : undefined
+      const routing = outOfBandRecord.reusable ? await this.routingService.getRouting(agentContext) : undefined
 
       const message = await this.didExchangeProtocol.createResponse(
-        messageContext.agentContext,
+        agentContext,
         connectionRecord,
         outOfBandRecord,
         routing
       )
       return new OutboundMessageContext(message, {
-        agentContext: messageContext.agentContext,
+        agentContext,
         connection: connectionRecord,
         outOfBand: outOfBandRecord,
       })

packages/core/src/modules/routing/__tests__/pickup.test.ts

@@ -78,6 +78,13 @@ describe('E2E Pick Up protocol', () => {
 
     mediatorRecipientConnection = await mediatorAgent.connections.returnWhenIsConnected(mediatorRecipientConnection!.id)
 
+    // Now they are connected, reinitialize recipient agent in order to lose the session (as with SubjectTransport it remains open)
+    await recipientAgent.shutdown()
+
+    recipientAgent = new Agent(recipientOptions)
+    recipientAgent.registerOutboundTransport(new SubjectOutboundTransport(subjectMap))
+    await recipientAgent.initialize()
+
     const message = 'hello pickup V1'
     await mediatorAgent.basicMessages.sendMessage(mediatorRecipientConnection.id, message)