Better handling of failed message migration attempts

Author: trevor-signal

ts/test-node/types/Message2_test.ts

@@ -270,15 +270,19 @@ describe('Message', () => {
           upgrade: v3,
         });
 
-        const context = getDefaultContext({ logger });
-        const upgradeSchema = async (message: MessageAttributesType) =>
-          toVersion3(
-            await toVersion2(await toVersion1(message, context), context),
-            context
-          );
-
-        const actual = await upgradeSchema(input);
+        const actual = await Message.upgradeSchema(input, getDefaultContext(), {
+          versions: [toVersion1, toVersion2, toVersion3],
+        });
         assert.deepEqual(actual, expected);
+
+        // if we try to upgrade it again, it will fail since it could not upgrade any
+        // versions
+        const upgradeAgainPromise = Message.upgradeSchema(
+          actual,
+          getDefaultContext(),
+          { versions: [toVersion1, toVersion2, toVersion3] }
+        );
+        await assert.isRejected(upgradeAgainPromise);
       });
 
       it('should skip out-of-order upgrade steps', async () => {
@@ -391,28 +395,25 @@ describe('Message', () => {
       assert.deepEqual(actual, expected);
     });
 
-    it('should return original message if upgrade function throws', async () => {
+    it('should throw if upgrade function throws', async () => {
       const upgrade = async () => {
         throw new Error('boom!');
       };
       const upgradeWithVersion = Message._withSchemaVersion({
-        schemaVersion: 3,
+        schemaVersion: 1,
         upgrade,
       });
 
       const input = getDefaultMessage({
         id: 'guid-guid-guid-guid',
         schemaVersion: 0,
       });
-      const expected = getDefaultMessage({
-        id: 'guid-guid-guid-guid',
-        schemaVersion: 0,
-      });
-      const actual = await upgradeWithVersion(
+
+      const upgradePromise = upgradeWithVersion(
         input,
         getDefaultContext({ logger })
       );
-      assert.deepEqual(actual, expected);
+      await assert.isRejected(upgradePromise);
     });
 
     it('should return original message if upgrade function returns null', async () => {
@@ -866,31 +867,5 @@ describe('Message', () => {
 
       assert.deepEqual({ ...message, schemaVersion: 14 }, result);
     });
-    it('if file does exist, but migration errors, does not increment version', async () => {
-      const message = getDefaultMessage({
-        schemaVersion: 13,
-        schemaMigrationAttempts: 0,
-        attachments: [
-          {
-            size: 128,
-            contentType: MIME.IMAGE_BMP,
-            path: 'no/file/here.png',
-            iv: 'iv',
-            digest: 'digest',
-            key: 'key',
-          },
-        ],
-        contact: [],
-      });
-      const result = await Message.upgradeSchema(message, {
-        ...getDefaultContext(),
-        doesAttachmentExist: async () => true,
-        ensureAttachmentIsReencryptable: async () => {
-          throw new Error("Can't reencrypt!");
-        },
-      });
-
-      assert.deepEqual(message, result);
-    });
   });
 });

ts/types/Attachment.ts

@@ -33,6 +33,7 @@ import { DAY } from '../util/durations';
 import { getMessageQueueTime } from '../util/getMessageQueueTime';
 import { getLocalAttachmentUrl } from '../util/getLocalAttachmentUrl';
 import type { ReencryptionInfo } from '../AttachmentCrypto';
+import { redactGenericText } from '../util/privacy';
 
 const MAX_WIDTH = 300;
 const MAX_HEIGHT = MAX_WIDTH * 1.5;
@@ -1238,3 +1239,11 @@ export function isAttachmentLocallySaved(
 ): attachment is LocallySavedAttachment {
   return Boolean(attachment.path);
 }
+
+export function getAttachmentIdForLogging(attachment: AttachmentType): string {
+  const { digest } = attachment;
+  if (typeof digest === 'string') {
+    return redactGenericText(digest);
+  }
+  return '[MissingDigest]';
+}

ts/types/Message2.ts

@@ -15,6 +15,7 @@ import type {
 } from './Attachment';
 import {
   captureDimensionsAndScreenshot,
+  getAttachmentIdForLogging,
   isAttachmentLocallySaved,
   removeSchemaVersion,
   replaceUnicodeOrderOverrides,
@@ -51,7 +52,6 @@ import { encryptLegacyAttachment } from '../util/encryptLegacyAttachment';
 import { deepClone } from '../util/deepClone';
 import { LONG_ATTACHMENT_LIMIT } from './Message';
 import * as Bytes from '../Bytes';
-import { redactGenericText } from '../util/privacy';
 
 export const GROUP = 'group';
 export const PRIVATE = 'private';
@@ -249,10 +249,11 @@ export const _withSchemaVersion = ({
       upgradedMessage = await upgrade(message, context);
     } catch (error) {
       logger.error(
-        `Message._withSchemaVersion: error updating message ${message.id}:`,
+        `Message._withSchemaVersion: error updating message ${message.id}, 
+        attempt ${message.schemaMigrationAttempts}:`,
         Errors.toLogFormat(error)
       );
-      return message;
+      throw error;
     }
 
     if (!isValid(upgradedMessage)) {
@@ -646,16 +647,16 @@ const toVersion14 = _withSchemaVersion({
       attachment,
       { logger, ensureAttachmentIsReencryptable, doesAttachmentExist }
     ) => {
-      const logId = `Message2.toVersion14(digest=${redactGenericText(attachment.digest ?? '')})`;
-
       if (!isAttachmentLocallySaved(attachment)) {
         return attachment;
       }
 
       if (!(await doesAttachmentExist(attachment.path))) {
         // Attachments may be missing, e.g. for quote thumbnails that reference messages
         // which have been deleted
-        logger.info(`${logId}: File does not exist`);
+        logger.info(
+          `Message2.toVersion14(id=${getAttachmentIdForLogging(attachment)}: File does not exist`
+        );
         return attachment;
       }
 
@@ -711,8 +712,17 @@ export const upgradeSchema = async (
     deleteOnDisk,
     logger,
     maxVersion = CURRENT_SCHEMA_VERSION,
-  }: ContextType
+  }: ContextType,
+  upgradeOptions: {
+    versions: ReadonlyArray<
+      (
+        message: MessageAttributesType,
+        context: ContextType
+      ) => Promise<MessageAttributesType>
+    >;
+  } = { versions: VERSIONS }
 ): Promise<MessageAttributesType> => {
+  const { versions } = upgradeOptions;
   if (!isFunction(readAttachmentData)) {
     throw new TypeError('context.readAttachmentData is required');
   }
@@ -748,30 +758,45 @@ export const upgradeSchema = async (
   }
 
   let message = rawMessage;
-  for (let index = 0, max = VERSIONS.length; index < max; index += 1) {
+  const startingVersion = message.schemaVersion ?? 0;
+  for (let index = 0, max = versions.length; index < max; index += 1) {
     if (maxVersion < index) {
       break;
     }
 
-    const currentVersion = VERSIONS[index];
-    // We really do want this intra-loop await because this is a chained async action,
-    //   each step dependent on the previous
-    // eslint-disable-next-line no-await-in-loop
-    message = await currentVersion(message, {
-      readAttachmentData,
-      writeNewAttachmentData,
-      makeObjectUrl,
-      revokeObjectUrl,
-      doesAttachmentExist,
-      ensureAttachmentIsReencryptable,
-      getImageDimensions,
-      makeImageThumbnail,
-      makeVideoScreenshot,
-      logger,
-      getRegionCode,
-      writeNewStickerData,
-      deleteOnDisk,
-    });
+    const currentVersion = versions[index];
+    try {
+      // We really do want this intra-loop await because this is a chained async action,
+      //   each step dependent on the previous
+      // eslint-disable-next-line no-await-in-loop
+      message = await currentVersion(message, {
+        readAttachmentData,
+        writeNewAttachmentData,
+        makeObjectUrl,
+        revokeObjectUrl,
+        doesAttachmentExist,
+        ensureAttachmentIsReencryptable,
+        getImageDimensions,
+        makeImageThumbnail,
+        makeVideoScreenshot,
+        logger,
+        getRegionCode,
+        writeNewStickerData,
+        deleteOnDisk,
+      });
+    } catch (e) {
+      // Throw the error if we were unable to upgrade the message at all
+      if (message.schemaVersion === startingVersion) {
+        throw e;
+      } else {
+        // Otherwise, return the message upgraded as far as it could be. On the next
+        // migration attempt, it will fail.
+        logger.error(
+          `Upgraded message from ${startingVersion} -> ${message.schemaVersion}; failed on upgrade to ${index}`
+        );
+        break;
+      }
+    }
   }
 
   return message;

ts/util/downloadAttachment.ts

@@ -5,11 +5,11 @@ import {
   type AttachmentType,
   mightBeOnBackupTier,
   AttachmentVariant,
+  getAttachmentIdForLogging,
 } from '../types/Attachment';
 import { downloadAttachment as doDownloadAttachment } from '../textsecure/downloadAttachment';
 import { MediaTier } from '../types/AttachmentDownload';
 import * as log from '../logging/log';
-import { redactGenericText } from './privacy';
 import { HTTPError } from '../textsecure/Errors';
 import { toLogFormat } from '../types/errors';
 import type { ReencryptedAttachmentV2 } from '../AttachmentCrypto';
@@ -25,10 +25,10 @@ export async function downloadAttachment({
   variant?: AttachmentVariant;
   dependencies?: { downloadAttachmentFromServer: typeof doDownloadAttachment };
 }): Promise<ReencryptedAttachmentV2> {
-  const redactedDigest = redactGenericText(attachment.digest ?? '');
+  const attachmentId = getAttachmentIdForLogging(attachment);
   const variantForLogging =
     variant !== AttachmentVariant.Default ? `[${variant}]` : '';
-  const dataId = `${redactedDigest}${variantForLogging}`;
+  const dataId = `${attachmentId}${variantForLogging}`;
   const logId = `downloadAttachmentUtil(${dataId})`;
 
   const { server } = window.textsecure;

ts/util/ensureAttachmentIsReencryptable.ts

@@ -17,11 +17,11 @@ import {
   hasAllOriginalEncryptionInfo,
   isReencryptableToSameDigest,
   isReencryptableWithNewEncryptionInfo,
+  getAttachmentIdForLogging,
 } from '../types/Attachment';
 import { strictAssert } from './assert';
 import * as logging from '../logging/log';
 import { fromBase64, toBase64 } from '../Bytes';
-import { redactGenericText } from './privacy';
 import { toLogFormat } from '../types/errors';
 
 /**
@@ -37,7 +37,6 @@ import { toLogFormat } from '../types/errors';
 export async function ensureAttachmentIsReencryptable(
   attachment: LocallySavedAttachment
 ): Promise<ReencryptableAttachment> {
-  const logId = `ensureAttachmentIsReencryptable(digest=${redactGenericText(attachment.digest ?? '')})`;
   if (isReencryptableToSameDigest(attachment)) {
     return attachment;
   }
@@ -54,6 +53,8 @@ export async function ensureAttachmentIsReencryptable(
         isReencryptableToSameDigest: true,
       };
     } catch (e) {
+      const logId = `ensureAttachmentIsReencryptable(digest=${getAttachmentIdForLogging(attachment)})`;
+
       if (e instanceof ReencryptedDigestMismatchError) {
         logging.info(
           `${logId}: Unable to reencrypt attachment to original digest; must have had non-zero padding`