1 files changed, 16 insertions, 15 deletions
diff --git a/deno/mail-relay/aws/mail.ts b/deno/mail-relay/aws/mail.ts
index cc05d23..0f86705 100644
--- a/deno/mail-relay/aws/mail.ts
+++ b/deno/mail-relay/aws/mail.ts
@@ -8,42 +8,43 @@ export class AwsMailMessageIdRewriteHook implements MailDeliverHook {
   }
 
   async callback(context: MailDeliverContext): Promise<void> {
-    console.info("Rewrite message ids...");
     const addresses = context.mail.simpleFindAllAddresses();
-    console.info(`Addresses found in mail: ${addresses.join(", ")}.`);
     for (const address of addresses) {
       const awsMessageId = await this.#lookup(address);
       if (awsMessageId != null && awsMessageId.length !== 0) {
-        console.info(`Rewrite ${address} to ${awsMessageId}.`);
+        console.info(`Rewrite address-line string in mail: ${address} => ${awsMessageId}.`);
         context.mail.raw = context.mail.raw.replaceAll(address, awsMessageId);
       }
     }
-    console.info("Done rewrite message ids.");
   }
 }
 
 export class AwsMailMessageIdSaveHook implements MailDeliverHook {
   readonly #record;
 
-  constructor(record: (original: string, aws: string) => Promise<void>) {
+  constructor(
+    record: (
+      original: string,
+      aws: string,
+      context: MailDeliverContext,
+    ) => Promise<void>,
+  ) {
     this.#record = record;
   }
 
   async callback(context: MailDeliverContext): Promise<void> {
-    console.info("Save aws message ids...");
-    const messageId = context.mail
-      .startSimpleParse()
-      .sections()
-      .headers()
-      .messageId();
+    const messageId = context.mail.parsed.messageId;
     if (messageId == null) {
-      console.info("Original mail does not have message id. Skip saving.");
+      console.warn("Original mail does not have message id, skip saving message id map.");
       return;
     }
     if (context.result.awsMessageId != null) {
-      console.info(`Saving ${messageId} => ${context.result.awsMessageId}.`);
-      await this.#record(messageId, context.result.awsMessageId);
+      console.info(`Save message id map: ${messageId} => ${context.result.awsMessageId}.`);
+      context.mail.raw = context.mail.raw.replaceAll(
+        messageId,
+        context.result.awsMessageId,
+      );
+      await this.#record(messageId, context.result.awsMessageId, context);
     }
-    console.info("Done save message ids.");
   }
 }