Self-repairing message counter

libtextsecure/test/_test.js

@@ -77,4 +77,6 @@ before(async () => {
       err && err.stack ? err.stack : err
     );
   }
+
+  await window.Signal.Util.initializeMessageCounter();
 });

test/_test.js

@@ -88,6 +88,7 @@ before(async () => {
       err && err.stack ? err.stack : err
     );
   }
+  await window.Signal.Util.initializeMessageCounter();
   await window.Signal.Data.removeAll();
   await window.storage.fetch();
 });

ts/background.ts

@@ -139,6 +139,8 @@ export async function startApp(): Promise<void> {
     );
   }
 
+  await window.Signal.Util.initializeMessageCounter();
+
   // Initialize WebAPI as early as possible
   let server: WebAPIType | undefined;
   let messageReceiver: MessageReceiver | undefined;

ts/sql/Client.ts

@@ -277,6 +277,8 @@ const dataInterface: ClientInterface = {
   processGroupCallRingCancelation,
   cleanExpiredGroupCallRings,
 
+  getMaxMessageCounter,
+
   getStatisticsForLogging,
 
   // Test-only
@@ -1656,6 +1658,10 @@ async function updateAllConversationColors(
   );
 }
 
+function getMaxMessageCounter(): Promise<number | undefined> {
+  return channels.getMaxMessageCounter();
+}
+
 function getStatisticsForLogging(): Promise<Record<string, string>> {
   return channels.getStatisticsForLogging();
 }

ts/sql/Interface.ts

@@ -486,6 +486,7 @@ export type DataInterface = {
     }
   ) => Promise<void>;
 
+  getMaxMessageCounter(): Promise<number | undefined>;
   getStatisticsForLogging(): Promise<Record<string, string>>;
 };
 

ts/sql/Server.ts

@@ -268,6 +268,8 @@ const dataInterface: ServerInterface = {
   processGroupCallRingCancelation,
   cleanExpiredGroupCallRings,
 
+  getMaxMessageCounter,
+
   getStatisticsForLogging,
 
   // Server-only
@@ -6493,6 +6495,25 @@ async function cleanExpiredGroupCallRings(): Promise<void> {
   });
 }
 
+async function getMaxMessageCounter(): Promise<number | undefined> {
+  const db = getInstance();
+
+  return db
+    .prepare<EmptyQuery>(
+      `
+    SELECT MAX(counter)
+    FROM
+      (
+        SELECT MAX(received_at) AS counter FROM messages
+        UNION
+        SELECT MAX(timestamp) AS counter FROM unprocessed
+      )
+    `
+    )
+    .pluck()
+    .get();
+}
+
 async function getStatisticsForLogging(): Promise<Record<string, string>> {
   const counts = await pProps({
     messageCount: getMessageCount(),

ts/util/incrementMessageCounter.ts

@@ -1,16 +1,57 @@
 // Copyright 2021 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
-import { debounce } from 'lodash';
+import { debounce, isNumber } from 'lodash';
+
+import { strictAssert } from './assert';
+import Data from '../sql/Client';
 
 let receivedAtCounter: number | undefined;
 
-export function incrementMessageCounter(): number {
-  if (!receivedAtCounter) {
-    receivedAtCounter =
-      Number(localStorage.getItem('lastReceivedAtCounter')) || Date.now();
+export async function initializeMessageCounter(): Promise<void> {
+  strictAssert(
+    receivedAtCounter === undefined,
+    'incrementMessageCounter: already initialized'
+  );
+
+  const storedCounter = Number(localStorage.getItem('lastReceivedAtCounter'));
+  const dbCounter = await Data.getMaxMessageCounter();
+
+  if (isNumber(dbCounter) && isNumber(storedCounter)) {
+    window.log.info(
+      'initializeMessageCounter: picking max of db/stored counters'
+    );
+    receivedAtCounter = Math.max(dbCounter, storedCounter);
+
+    if (receivedAtCounter !== storedCounter) {
+      window.log.warn(
+        'initializeMessageCounter: mismatch between db/stored counters'
+      );
+    }
+  } else if (isNumber(storedCounter)) {
+    window.log.info('initializeMessageCounter: picking stored counter');
+    receivedAtCounter = storedCounter;
+  } else if (isNumber(dbCounter)) {
+    window.log.info(
+      'initializeMessageCounter: picking fallback counter from the database'
+    );
+    receivedAtCounter = dbCounter;
+  } else {
+    window.log.info('initializeMessageCounter: defaulting to Date.now()');
+    receivedAtCounter = Date.now();
   }
 
+  if (storedCounter !== receivedAtCounter) {
+    localStorage.setItem('lastReceivedAtCounter', String(receivedAtCounter));
+  }
+}
+
+export function incrementMessageCounter(): number {
+  strictAssert(
+    receivedAtCounter !== undefined,
+    'incrementMessageCounter: not initialized'
+  );
+
   receivedAtCounter += 1;
   debouncedUpdateLastReceivedAt();
 

ts/util/index.ts

@@ -15,6 +15,7 @@ import { getTextWithMentions } from './getTextWithMentions';
 import { getUserAgent } from './getUserAgent';
 import { hasExpired } from './hasExpired';
 import {
+  initializeMessageCounter,
   incrementMessageCounter,
   flushMessageCounter,
 } from './incrementMessageCounter';
@@ -61,6 +62,7 @@ export {
   getUserAgent,
   hasExpired,
   incrementMessageCounter,
+  initializeMessageCounter,
   isFileDangerous,
   longRunningTaskWrapper,
   makeLookup,