Normalize message attachments

ts/sql/hydration.ts

@@ -1,10 +1,11 @@
 // Copyright 2025 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
+import { groupBy } from 'lodash';
 import type { ReadStatus } from '../messages/MessageReadStatus';
 import type { SeenStatus } from '../MessageSeenStatus';
 import type { ServiceIdString } from '../types/ServiceId';
-import { dropNull } from '../util/dropNull';
+import { dropNull, shallowDropNull } from '../util/dropNull';
 
 /* eslint-disable camelcase */
 
@@ -12,7 +13,23 @@ import type {
   MessageTypeUnhydrated,
   MessageType,
   MESSAGE_COLUMNS,
+  ReadableDB,
+  MessageAttachmentDBType,
 } from './Interface';
+import {
+  batchMultiVarQuery,
+  convertOptionalIntegerToBoolean,
+  jsonToObject,
+  sql,
+  sqlJoin,
+} from './util';
+import type { AttachmentType } from '../types/Attachment';
+import { IMAGE_JPEG, stringToMIMEType } from '../types/MIME';
+import { strictAssert } from '../util/assert';
+import { sqlLogger } from './sqlLogger';
+import type { MessageAttributesType } from '../model-types';
+
+export const ROOT_MESSAGE_ATTACHMENT_EDIT_HISTORY_INDEX = -1;
 
 function toBoolean(value: number | null): boolean | undefined {
   if (value == null) {
@@ -21,7 +38,27 @@ function toBoolean(value: number | null): boolean | undefined {
   return value === 1;
 }
 
-export function hydrateMessage(row: MessageTypeUnhydrated): MessageType {
+export function hydrateMessage(
+  db: ReadableDB,
+  row: MessageTypeUnhydrated
+): MessageType {
+  return hydrateMessages(db, [row])[0];
+}
+
+export function hydrateMessages(
+  db: ReadableDB,
+  unhydratedMessages: Array<MessageTypeUnhydrated>
+): Array<MessageType> {
+  const messagesWithColumnsHydrated = unhydratedMessages.map(
+    hydrateMessageTableColumns
+  );
+
+  return hydrateMessagesWithAttachments(db, messagesWithColumnsHydrated);
+}
+
+export function hydrateMessageTableColumns(
+  row: MessageTypeUnhydrated
+): MessageType {
   const {
     json,
     id,
@@ -29,9 +66,6 @@ export function hydrateMessage(row: MessageTypeUnhydrated): MessageType {
     conversationId,
     expirationStartTimestamp,
     expireTimer,
-    hasAttachments,
-    hasFileAttachments,
-    hasVisualMediaAttachments,
     isErased,
     isViewOnce,
     mentionsMe,
@@ -63,9 +97,6 @@ export function hydrateMessage(row: MessageTypeUnhydrated): MessageType {
     conversationId: conversationId || '',
     expirationStartTimestamp: dropNull(expirationStartTimestamp),
     expireTimer: dropNull(expireTimer) as MessageType['expireTimer'],
-    hasAttachments: toBoolean(hasAttachments),
-    hasFileAttachments: toBoolean(hasFileAttachments),
-    hasVisualMediaAttachments: toBoolean(hasVisualMediaAttachments),
     isErased: toBoolean(isErased),
     isViewOnce: toBoolean(isViewOnce),
     mentionsMe: toBoolean(mentionsMe),
@@ -86,3 +117,299 @@ export function hydrateMessage(row: MessageTypeUnhydrated): MessageType {
     unidentifiedDeliveryReceived: toBoolean(unidentifiedDeliveryReceived),
   };
 }
+
+export function getAttachmentReferencesForMessages(
+  db: ReadableDB,
+  messageIds: Array<string>
+): Array<MessageAttachmentDBType> {
+  return batchMultiVarQuery(
+    db,
+    messageIds,
+    (
+      messageIdBatch: ReadonlyArray<string>,
+      persistent: boolean
+    ): Array<MessageAttachmentDBType> => {
+      const [query, params] = sql`
+      SELECT * FROM message_attachments 
+      WHERE messageId IN (${sqlJoin(messageIdBatch)});
+    `;
+
+      return db
+        .prepare(query, { persistent })
+        .all<MessageAttachmentDBType>(params);
+    }
+  );
+}
+
+function hydrateMessagesWithAttachments(
+  db: ReadableDB,
+  messagesWithoutAttachments: Array<MessageType>
+): Array<MessageType> {
+  const attachmentReferencesForAllMessages = getAttachmentReferencesForMessages(
+    db,
+    messagesWithoutAttachments.map(msg => msg.id)
+  );
+  const attachmentReferencesByMessage = groupBy(
+    attachmentReferencesForAllMessages,
+    'messageId'
+  );
+
+  return messagesWithoutAttachments.map(msg => {
+    const attachmentReferences = attachmentReferencesByMessage[msg.id] ?? [];
+    if (!attachmentReferences.length) {
+      return msg;
+    }
+
+    const attachmentsByEditHistoryIndex = groupBy(
+      attachmentReferences,
+      'editHistoryIndex'
+    );
+
+    const message = hydrateMessageRootOrRevisionWithAttachments(
+      msg,
+      attachmentsByEditHistoryIndex[
+        ROOT_MESSAGE_ATTACHMENT_EDIT_HISTORY_INDEX
+      ] ?? []
+    );
+
+    if (message.editHistory) {
+      message.editHistory = message.editHistory.map((editHistory, idx) => {
+        return hydrateMessageRootOrRevisionWithAttachments(
+          editHistory,
+          attachmentsByEditHistoryIndex[idx] ?? []
+        );
+      });
+    }
+
+    return message;
+  });
+}
+
+function hydrateMessageRootOrRevisionWithAttachments<
+  T extends Pick<
+    MessageAttributesType,
+    | 'attachments'
+    | 'bodyAttachment'
+    | 'contact'
+    | 'preview'
+    | 'quote'
+    | 'sticker'
+  >,
+>(message: T, messageAttachments: Array<MessageAttachmentDBType>): T {
+  const attachmentsByType = groupBy(
+    messageAttachments,
+    'attachmentType'
+  ) as Record<
+    MessageAttachmentDBType['attachmentType'],
+    Array<MessageAttachmentDBType>
+  >;
+
+  const standardAttachments = attachmentsByType.attachment ?? [];
+  const bodyAttachments = attachmentsByType['long-message'] ?? [];
+  const quoteAttachments = attachmentsByType.quote ?? [];
+  const previewAttachments = attachmentsByType.preview ?? [];
+  const contactAttachments = attachmentsByType.contact ?? [];
+  const stickerAttachment = (attachmentsByType.sticker ?? []).find(
+    sticker => sticker.orderInMessage === 0
+  );
+
+  const hydratedMessage = structuredClone(message);
+
+  if (standardAttachments.length) {
+    hydratedMessage.attachments = standardAttachments
+      .sort((a, b) => a.orderInMessage - b.orderInMessage)
+      .map(convertAttachmentDBFieldsToAttachmentType);
+  }
+
+  if (bodyAttachments[0]) {
+    hydratedMessage.bodyAttachment = convertAttachmentDBFieldsToAttachmentType(
+      bodyAttachments[0]
+    );
+  }
+
+  hydratedMessage.quote?.attachments.forEach((quoteAttachment, idx) => {
+    const quoteThumbnail = quoteAttachments.find(
+      attachment => attachment.orderInMessage === idx
+    );
+    if (quoteThumbnail) {
+      // eslint-disable-next-line no-param-reassign
+      quoteAttachment.thumbnail =
+        convertAttachmentDBFieldsToAttachmentType(quoteThumbnail);
+    }
+  });
+
+  hydratedMessage.preview?.forEach((preview, idx) => {
+    const previewAttachment = previewAttachments.find(
+      attachment => attachment.orderInMessage === idx
+    );
+
+    if (previewAttachment) {
+      // eslint-disable-next-line no-param-reassign
+      preview.image =
+        convertAttachmentDBFieldsToAttachmentType(previewAttachment);
+    }
+  });
+
+  hydratedMessage.contact?.forEach((contact, idx) => {
+    const contactAttachment = contactAttachments.find(
+      attachment => attachment.orderInMessage === idx
+    );
+    if (contactAttachment && contact.avatar) {
+      // eslint-disable-next-line no-param-reassign
+      contact.avatar.avatar =
+        convertAttachmentDBFieldsToAttachmentType(contactAttachment);
+    }
+  });
+
+  if (hydratedMessage.sticker && stickerAttachment) {
+    hydratedMessage.sticker.data =
+      convertAttachmentDBFieldsToAttachmentType(stickerAttachment);
+  }
+
+  return hydratedMessage;
+}
+
+function convertAttachmentDBFieldsToAttachmentType(
+  dbFields: MessageAttachmentDBType
+): AttachmentType {
+  const messageAttachment = shallowDropNull(dbFields);
+  strictAssert(messageAttachment != null, 'must exist');
+
+  const {
+    clientUuid,
+    size,
+    contentType,
+    plaintextHash,
+    path,
+    localKey,
+    caption,
+    blurHash,
+    height,
+    width,
+    digest,
+    iv,
+    key,
+    downloadPath,
+    flags,
+    fileName,
+    version,
+    incrementalMac,
+    incrementalMacChunkSize: chunkSize,
+    transitCdnKey: cdnKey,
+    transitCdnNumber: cdnNumber,
+    transitCdnUploadTimestamp: uploadTimestamp,
+    error,
+    pending,
+    wasTooBig,
+    isCorrupted,
+    backfillError,
+    storyTextAttachmentJson,
+    copiedFromQuotedAttachment,
+    isReencryptableToSameDigest,
+    localBackupPath,
+  } = messageAttachment;
+
+  const result: AttachmentType = {
+    clientUuid,
+    size,
+    contentType: stringToMIMEType(contentType),
+    plaintextHash,
+    path,
+    localKey,
+    caption,
+    blurHash,
+    height,
+    width,
+    digest,
+    iv,
+    key,
+    downloadPath,
+    localBackupPath,
+    flags,
+    fileName,
+    version,
+    incrementalMac,
+    chunkSize,
+    cdnKey,
+    cdnNumber,
+    uploadTimestamp,
+    pending: convertOptionalIntegerToBoolean(pending),
+    error: convertOptionalIntegerToBoolean(error),
+    wasTooBig: convertOptionalIntegerToBoolean(wasTooBig),
+    copied: convertOptionalIntegerToBoolean(copiedFromQuotedAttachment),
+    isCorrupted: convertOptionalIntegerToBoolean(isCorrupted),
+    backfillError: convertOptionalIntegerToBoolean(backfillError),
+    isReencryptableToSameDigest: convertOptionalIntegerToBoolean(
+      isReencryptableToSameDigest
+    ),
+    textAttachment: storyTextAttachmentJson
+      ? jsonToObject(storyTextAttachmentJson)
+      : undefined,
+    ...(messageAttachment.backupMediaName
+      ? {
+          backupLocator: {
+            mediaName: messageAttachment.backupMediaName,
+            cdnNumber: messageAttachment.backupCdnNumber,
+          },
+        }
+      : {}),
+    ...(messageAttachment.thumbnailPath
+      ? {
+          thumbnail: {
+            path: messageAttachment.thumbnailPath,
+            size: messageAttachment.thumbnailSize ?? 0,
+            contentType: messageAttachment.thumbnailContentType
+              ? stringToMIMEType(messageAttachment.thumbnailContentType)
+              : IMAGE_JPEG,
+            localKey: messageAttachment.thumbnailLocalKey,
+            version: messageAttachment.thumbnailVersion,
+          },
+        }
+      : {}),
+    ...(messageAttachment.screenshotPath
+      ? {
+          screenshot: {
+            path: messageAttachment.screenshotPath,
+            size: messageAttachment.screenshotSize ?? 0,
+            contentType: messageAttachment.screenshotContentType
+              ? stringToMIMEType(messageAttachment.screenshotContentType)
+              : IMAGE_JPEG,
+            localKey: messageAttachment.screenshotLocalKey,
+            version: messageAttachment.screenshotVersion,
+          },
+        }
+      : {}),
+    ...(messageAttachment.backupThumbnailPath
+      ? {
+          thumbnailFromBackup: {
+            path: messageAttachment.backupThumbnailPath,
+            size: messageAttachment.backupThumbnailSize ?? 0,
+            contentType: messageAttachment.backupThumbnailContentType
+              ? stringToMIMEType(messageAttachment.backupThumbnailContentType)
+              : IMAGE_JPEG,
+            localKey: messageAttachment.backupThumbnailLocalKey,
+            version: messageAttachment.backupThumbnailVersion,
+          },
+        }
+      : {}),
+  };
+
+  if (result.isReencryptableToSameDigest === false) {
+    if (
+      !messageAttachment.reencryptionIv ||
+      !messageAttachment.reencryptionKey ||
+      !messageAttachment.reencryptionDigest
+    ) {
+      sqlLogger.warn(
+        'Attachment missing reencryption info despite not being reencryptable'
+      );
+      return result;
+    }
+    result.reencryptionInfo = {
+      iv: messageAttachment.reencryptionIv,
+      key: messageAttachment.reencryptionKey,
+      digest: messageAttachment.reencryptionDigest,
+    };
+  }
+  return result;
+}