refactor(core): separate logic for inbound plaintext and encrypted message (openwallet-foundation#581)

* rename wire message to encrypted message

Signed-off-by: Jakub Koci <[email protected]>

Author: jakubkoci

Diff for: packages/core/src/agent/Agent.ts

@@ -204,8 +204,8 @@ export class Agent {
     return this.wallet.publicDid
   }
 
-  public async receiveMessage(inboundPackedMessage: unknown, session?: TransportSession) {
-    return await this.messageReceiver.receiveMessage(inboundPackedMessage, session)
+  public async receiveMessage(inboundMessage: unknown, session?: TransportSession) {
+    return await this.messageReceiver.receiveMessage(inboundMessage, session)
   }
 
   public get injectionContainer() {

Diff for: packages/core/src/agent/EnvelopeService.ts

@@ -1,5 +1,5 @@
 import type { Logger } from '../logger'
-import type { UnpackedMessageContext, WireMessage } from '../types'
+import type { DecryptedMessageContext, EncryptedMessage } from '../types'
 import type { AgentMessage } from './AgentMessage'
 
 import { inject, scoped, Lifecycle } from 'tsyringe'
@@ -28,7 +28,7 @@ class EnvelopeService {
     this.config = agentConfig
   }
 
-  public async packMessage(payload: AgentMessage, keys: EnvelopeKeys): Promise<WireMessage> {
+  public async packMessage(payload: AgentMessage, keys: EnvelopeKeys): Promise<EncryptedMessage> {
     const { routingKeys, senderKey } = keys
     let recipientKeys = keys.recipientKeys
 
@@ -37,29 +37,29 @@ class EnvelopeService {
 
     this.logger.debug(`Pack outbound message ${message['@type']}`)
 
-    let wireMessage = await this.wallet.pack(message, recipientKeys, senderKey ?? undefined)
+    let encryptedMessage = await this.wallet.pack(message, recipientKeys, senderKey ?? undefined)
 
     // If the message has routing keys (mediator) pack for each mediator
     for (const routingKey of routingKeys) {
       const forwardMessage = new ForwardMessage({
         // Forward to first recipient key
         to: recipientKeys[0],
-        message: wireMessage,
+        message: encryptedMessage,
       })
       recipientKeys = [routingKey]
       this.logger.debug('Forward message created', forwardMessage)
 
       const forwardJson = forwardMessage.toJSON({ useLegacyDidSovPrefix: this.config.useLegacyDidSovPrefix })
 
       // Forward messages are anon packed
-      wireMessage = await this.wallet.pack(forwardJson, [routingKey], undefined)
+      encryptedMessage = await this.wallet.pack(forwardJson, [routingKey], undefined)
     }
 
-    return wireMessage
+    return encryptedMessage
   }
 
-  public async unpackMessage(packedMessage: WireMessage): Promise<UnpackedMessageContext> {
-    return this.wallet.unpack(packedMessage)
+  public async unpackMessage(encryptedMessage: EncryptedMessage): Promise<DecryptedMessageContext> {
+    return this.wallet.unpack(encryptedMessage)
   }
 }
 

Diff for: packages/core/src/agent/MessageReceiver.ts

@@ -1,7 +1,7 @@
 import type { Logger } from '../logger'
 import type { ConnectionRecord } from '../modules/connections'
 import type { InboundTransport } from '../transport'
-import type { UnpackedMessageContext, UnpackedMessage, WireMessage } from '../types'
+import type { DecryptedMessageContext, PlaintextMessage, EncryptedMessage } from '../types'
 import type { AgentMessage } from './AgentMessage'
 import type { TransportSession } from './TransportService'
 
@@ -59,21 +59,29 @@ export class MessageReceiver {
   }
 
   /**
-   * Receive and handle an inbound DIDComm message. It will unpack the message, transform it
+   * Receive and handle an inbound DIDComm message. It will decrypt the message, transform it
    * to it's corresponding message class and finally dispatch it to the dispatcher.
    *
-   * @param inboundPackedMessage the message to receive and handle
+   * @param inboundMessage the message to receive and handle
    */
-  public async receiveMessage(inboundPackedMessage: unknown, session?: TransportSession) {
-    if (typeof inboundPackedMessage !== 'object' || inboundPackedMessage == null) {
-      throw new AriesFrameworkError('Invalid message received. Message should be object')
+  public async receiveMessage(inboundMessage: unknown, session?: TransportSession) {
+    this.logger.debug(`Agent ${this.config.label} received message`)
+
+    if (this.isPlaintextMessage(inboundMessage)) {
+      await this.receivePlaintextMessage(inboundMessage)
+    } else {
+      await this.receiveEncryptedMessage(inboundMessage as EncryptedMessage, session)
     }
+  }
 
-    this.logger.debug(`Agent ${this.config.label} received message`)
+  private async receivePlaintextMessage(plaintextMessage: PlaintextMessage) {
+    const message = await this.transformAndValidate(plaintextMessage)
+    const messageContext = new InboundMessageContext(message, {})
+    await this.dispatcher.dispatch(messageContext)
+  }
 
-    const unpackedMessage = await this.unpackMessage(inboundPackedMessage as WireMessage)
-    const senderKey = unpackedMessage.senderVerkey
-    const recipientKey = unpackedMessage.recipientVerkey
+  private async receiveEncryptedMessage(encryptedMessage: EncryptedMessage, session?: TransportSession) {
+    const { plaintextMessage, senderKey, recipientKey } = await this.decryptMessage(encryptedMessage)
 
     let connection: ConnectionRecord | null = null
 
@@ -90,27 +98,11 @@ export class MessageReceiver {
     }
 
     this.logger.info(
-      `Received message with type '${unpackedMessage.message['@type']}' from connection ${connection?.id} (${connection?.theirLabel})`,
-      unpackedMessage.message
+      `Received message with type '${plaintextMessage['@type']}' from connection ${connection?.id} (${connection?.theirLabel})`,
+      plaintextMessage
     )
 
-    let message: AgentMessage | null = null
-    try {
-      message = await this.transformMessage(unpackedMessage)
-      await this.validateMessage(message)
-    } catch (error) {
-      if (connection) await this.sendProblemReportMessage(error.message, connection, unpackedMessage)
-      throw error
-    }
-
-    const messageContext = new InboundMessageContext(message, {
-      // Only make the connection available in message context if the connection is ready
-      // To prevent unwanted usage of unready connections. Connections can still be retrieved from
-      // Storage if the specific protocol allows an unready connection to be used.
-      connection: connection?.isReady ? connection : undefined,
-      senderVerkey: senderKey,
-      recipientVerkey: recipientKey,
-    })
+    const message = await this.transformAndValidate(plaintextMessage, connection)
 
     // We want to save a session if there is a chance of returning outbound message via inbound transport.
     // That can happen when inbound message has `return_route` set to `all` or `thread`.
@@ -131,53 +123,68 @@ export class MessageReceiver {
       this.transportService.saveSession(session)
     }
 
+    const messageContext = new InboundMessageContext(message, {
+      // Only make the connection available in message context if the connection is ready
+      // To prevent unwanted usage of unready connections. Connections can still be retrieved from
+      // Storage if the specific protocol allows an unready connection to be used.
+      connection: connection?.isReady ? connection : undefined,
+      senderVerkey: senderKey,
+      recipientVerkey: recipientKey,
+    })
     await this.dispatcher.dispatch(messageContext)
   }
 
   /**
-   * Unpack a message using the envelope service.
-   * If message is not packed, it will be returned as is, but in the unpacked message structure
+   * Decrypt a message using the envelope service.
    *
-   * @param packedMessage the received, probably packed, message to unpack
+   * @param message the received inbound message to decrypt
    */
-  private async unpackMessage(packedMessage: WireMessage): Promise<UnpackedMessageContext> {
-    // If the inbound message has no @type field we assume
-    // the message is packed and must be unpacked first
-    if (!this.isUnpackedMessage(packedMessage)) {
-      try {
-        return await this.envelopeService.unpackMessage(packedMessage)
-      } catch (error) {
-        this.logger.error('error while unpacking message', {
-          error,
-          packedMessage,
-          errorMessage: error instanceof Error ? error.message : error,
-        })
-        throw error
-      }
+  private async decryptMessage(message: EncryptedMessage): Promise<DecryptedMessageContext> {
+    try {
+      return await this.envelopeService.unpackMessage(message)
+    } catch (error) {
+      this.logger.error('Error while decrypting message', {
+        error,
+        encryptedMessage: message,
+        errorMessage: error instanceof Error ? error.message : error,
+      })
+      throw error
     }
+  }
 
-    // If the message does have an @type field we assume
-    // the message is already unpacked an use it directly
-    else {
-      const unpackedMessage: UnpackedMessageContext = { message: packedMessage }
-      return unpackedMessage
+  private isPlaintextMessage(message: unknown): message is PlaintextMessage {
+    if (typeof message !== 'object' || message == null) {
+      throw new AriesFrameworkError('Invalid message received. Message should be object')
     }
+    // If the message does have an @type field we assume the message is in plaintext and it is not encrypted.
+    return '@type' in message
   }
 
-  private isUnpackedMessage(message: Record<string, unknown>): message is UnpackedMessage {
-    return '@type' in message
+  private async transformAndValidate(
+    plaintextMessage: PlaintextMessage,
+    connection?: ConnectionRecord | null
+  ): Promise<AgentMessage> {
+    let message: AgentMessage
+    try {
+      message = await this.transformMessage(plaintextMessage)
+      await this.validateMessage(message)
+    } catch (error) {
+      if (connection) await this.sendProblemReportMessage(error.message, connection, plaintextMessage)
+      throw error
+    }
+    return message
   }
 
   /**
-   * Transform an unpacked DIDComm message into it's corresponding message class. Will look at all message types in the registered handlers.
+   * Transform an plaintext DIDComm message into it's corresponding message class. Will look at all message types in the registered handlers.
    *
-   * @param unpackedMessage the unpacked message for which to transform the message in to a class instance
+   * @param message the plaintext message for which to transform the message in to a class instance
    */
-  private async transformMessage(unpackedMessage: UnpackedMessageContext): Promise<AgentMessage> {
+  private async transformMessage(message: PlaintextMessage): Promise<AgentMessage> {
     // replace did:sov:BzCbsNYhMrjHiqZDTUASHg;spec prefix for message type with https://didcomm.org
-    replaceLegacyDidSovPrefixOnMessage(unpackedMessage.message)
+    replaceLegacyDidSovPrefixOnMessage(message)
 
-    const messageType = unpackedMessage.message['@type']
+    const messageType = message['@type']
     const MessageClass = this.dispatcher.getMessageClassForType(messageType)
 
     if (!MessageClass) {
@@ -187,9 +194,7 @@ export class MessageReceiver {
     }
 
     // Cast the plain JSON object to specific instance of Message extended from AgentMessage
-    const message = JsonTransformer.fromJSON(unpackedMessage.message, MessageClass)
-
-    return message
+    return JsonTransformer.fromJSON(message, MessageClass)
   }
 
   /**
@@ -214,14 +219,14 @@ export class MessageReceiver {
    * Send the problem report message (https://didcomm.org/notification/1.0/problem-report) to the recipient.
    * @param message error message to send
    * @param connection connection to send the message to
-   * @param unpackedMessage received unpackedMessage
+   * @param plaintextMessage received inbound message
    */
   private async sendProblemReportMessage(
     message: string,
     connection: ConnectionRecord,
-    unpackedMessage: UnpackedMessageContext
+    plaintextMessage: PlaintextMessage
   ) {
-    if (unpackedMessage.message['@type'] === CommonMessageType.ProblemReport) {
+    if (plaintextMessage['@type'] === CommonMessageType.ProblemReport) {
       throw new AriesFrameworkError(message)
     }
     const problemReportMessage = new ProblemReportMessage({
@@ -231,7 +236,7 @@ export class MessageReceiver {
       },
     })
     problemReportMessage.setThread({
-      threadId: unpackedMessage.message['@id'],
+      threadId: plaintextMessage['@id'],
     })
     const outboundMessage = createOutboundMessage(connection, problemReportMessage)
     if (outboundMessage) {

Diff for: packages/core/src/agent/MessageSender.ts

@@ -1,6 +1,6 @@
 import type { DidCommService, ConnectionRecord } from '../modules/connections'
 import type { OutboundTransport } from '../transport/OutboundTransport'
-import type { OutboundMessage, OutboundPackage, WireMessage } from '../types'
+import type { OutboundMessage, OutboundPackage, EncryptedMessage } from '../types'
 import type { AgentMessage } from './AgentMessage'
 import type { EnvelopeKeys } from './EnvelopeService'
 import type { TransportSession } from './TransportService'
@@ -56,10 +56,10 @@ export class MessageSender {
     message: AgentMessage
     endpoint: string
   }): Promise<OutboundPackage> {
-    const wireMessage = await this.envelopeService.packMessage(message, keys)
+    const encryptedMessage = await this.envelopeService.packMessage(message, keys)
 
     return {
-      payload: wireMessage,
+      payload: encryptedMessage,
       responseRequested: message.hasAnyReturnRoute(),
       endpoint,
     }
@@ -72,18 +72,17 @@ export class MessageSender {
     if (!session.keys) {
       throw new AriesFrameworkError(`There are no keys for the given ${session.type} transport session.`)
     }
-    const wireMessage = await this.envelopeService.packMessage(message, session.keys)
-
-    await session.send(wireMessage)
+    const encryptedMessage = await this.envelopeService.packMessage(message, session.keys)
+    await session.send(encryptedMessage)
   }
 
   public async sendPackage({
     connection,
-    packedMessage,
+    encryptedMessage,
     options,
   }: {
     connection: ConnectionRecord
-    packedMessage: WireMessage
+    encryptedMessage: EncryptedMessage
     options?: { transportPriority?: TransportPriorityOptions }
   }) {
     const errors: Error[] = []
@@ -92,7 +91,7 @@ export class MessageSender {
     const session = this.transportService.findSessionByConnectionId(connection.id)
     if (session?.inboundMessage?.hasReturnRouting()) {
       try {
-        await session.send(packedMessage)
+        await session.send(encryptedMessage)
         return
       } catch (error) {
         errors.push(error)
@@ -114,7 +113,7 @@ export class MessageSender {
         for (const transport of this.outboundTransports) {
           if (transport.supportedSchemes.includes(service.protocolScheme)) {
             await transport.sendMessage({
-              payload: packedMessage,
+              payload: encryptedMessage,
               endpoint: service.serviceEndpoint,
               connectionId: connection.id,
             })
@@ -137,13 +136,13 @@ export class MessageSender {
     // If the other party shared a queue service endpoint in their did doc we queue the message
     if (queueService) {
       this.logger.debug(`Queue packed message for connection ${connection.id} (${connection.theirLabel})`)
-      this.messageRepository.add(connection.id, packedMessage)
+      this.messageRepository.add(connection.id, encryptedMessage)
       return
     }
 
     // Message is undeliverable
     this.logger.error(`Message is undeliverable to connection ${connection.id} (${connection.theirLabel})`, {
-      message: packedMessage,
+      message: encryptedMessage,
       errors,
       connection,
     })
@@ -217,8 +216,8 @@ export class MessageSender {
         senderKey: connection.verkey,
       }
 
-      const wireMessage = await this.envelopeService.packMessage(payload, keys)
-      this.messageRepository.add(connection.id, wireMessage)
+      const encryptedMessage = await this.envelopeService.packMessage(payload, keys)
+      this.messageRepository.add(connection.id, encryptedMessage)
       return
     }
 

Diff for: packages/core/src/agent/TransportService.ts

@@ -1,6 +1,6 @@
 import type { DidDoc, IndyAgentService } from '../modules/connections/models'
 import type { ConnectionRecord } from '../modules/connections/repository'
-import type { WireMessage } from '../types'
+import type { EncryptedMessage } from '../types'
 import type { AgentMessage } from './AgentMessage'
 import type { EnvelopeKeys } from './EnvelopeService'
 
@@ -64,5 +64,5 @@ export interface TransportSession {
   keys?: EnvelopeKeys
   inboundMessage?: AgentMessage
   connection?: ConnectionRecord
-  send(wireMessage: WireMessage): Promise<void>
+  send(encryptedMessage: EncryptedMessage): Promise<void>
 }