From afb432dceea26c53683b5d1b14cbbe68693fd588 Mon Sep 17 00:00:00 2001
From: Paul Schaub <vanitasvitae@fsfe.org>
Date: Wed, 10 Oct 2018 13:00:07 +0200
Subject: [PATCH] Add store methods for message counters

---
 .../smackx/omemo/CachingOmemoStore.java       | 25 +++++++++++++++++++
 .../smackx/omemo/FileBasedOmemoStore.java     | 24 ++++++++++++++++++
 .../jivesoftware/smackx/omemo/OmemoStore.java | 21 ++++++++++++++++
 .../smackx/omemo/OmemoStoreTest.java          |  7 ++++++
 4 files changed, 77 insertions(+)

diff --git a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/CachingOmemoStore.java b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/CachingOmemoStore.java
index ec4b23a46..7ab0bcde6 100644
--- a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/CachingOmemoStore.java
+++ b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/CachingOmemoStore.java
@@ -140,6 +140,30 @@ public class CachingOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Se
         }
     }
 
+    @Override
+    public void storeOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice, int counter) {
+        getCache(userDevice).messageCounters.put(contactsDevice, counter);
+        if (persistent != null) {
+            persistent.storeOmemoMessageCounter(userDevice, contactsDevice, counter);
+        }
+    }
+
+    @Override
+    public int loadOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        Integer counter = getCache(userDevice).messageCounters.get(contactsDevice);
+        if (counter == null && persistent != null) {
+            counter = persistent.loadOmemoMessageCounter(userDevice, contactsDevice);
+        }
+
+        if (counter == null) {
+            counter = 0;
+        }
+
+        getCache(userDevice).messageCounters.put(contactsDevice, counter);
+
+        return counter;
+    }
+
     @Override
     public void setDateOfLastReceivedMessage(OmemoDevice userDevice, OmemoDevice from, Date date) {
         getCache(userDevice).lastMessagesDates.put(from, date);
@@ -442,5 +466,6 @@ public class CachingOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_Se
         private final HashMap<OmemoDevice, Date> lastDeviceIdPublicationDates = new HashMap<>();
         private final HashMap<BareJid, OmemoCachedDeviceList> deviceLists = new HashMap<>();
         private Date lastRenewalDate = null;
+        private final HashMap<OmemoDevice, Integer> messageCounters = new HashMap<>();
     }
 }
diff --git a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/FileBasedOmemoStore.java b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/FileBasedOmemoStore.java
index 4ecefcda7..bab399c35 100644
--- a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/FileBasedOmemoStore.java
+++ b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/FileBasedOmemoStore.java
@@ -24,6 +24,7 @@ import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -353,6 +354,24 @@ public abstract class FileBasedOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigP
         return session.exists();
     }
 
+    @Override
+    public void storeOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice, int counter) {
+        File messageCounterFile = hierarchy.getDevicesMessageCounterPath(userDevice, contactsDevice);
+        writeIntegers(messageCounterFile, Collections.singleton(counter));
+    }
+
+    @Override
+    public int loadOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice) {
+        File messageCounterFile = hierarchy.getDevicesMessageCounterPath(userDevice, contactsDevice);
+        Set<Integer> integers = readIntegers(messageCounterFile);
+
+        if (integers == null || integers.isEmpty()) {
+            return 0;
+        }
+
+        return integers.iterator().next();
+    }
+
     @Override
     public OmemoCachedDeviceList loadCachedDeviceList(OmemoDevice userDevice, BareJid contact) {
         OmemoCachedDeviceList cachedDeviceList = new OmemoCachedDeviceList();
@@ -732,6 +751,7 @@ public abstract class FileBasedOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigP
         static final String SESSION = "session";
         static final String DEVICE_LIST_ACTIVE = "activeDevices";
         static final String DEVICE_LIST_INAVTIVE = "inactiveDevices";
+        static final String MESSAGE_COUNTER = "messageCounter";
 
         File basePath;
 
@@ -815,6 +835,10 @@ public abstract class FileBasedOmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigP
             return new File(getContactsDir(userDevice, contact), DEVICE_LIST_INAVTIVE);
         }
 
+        File getDevicesMessageCounterPath(OmemoDevice userDevice, OmemoDevice otherDevice) {
+            return new File(getContactsDir(userDevice, otherDevice), MESSAGE_COUNTER);
+        }
+
         private static File createFile(File f) throws IOException {
             File p = f.getParentFile();
             createDirectory(p);
diff --git a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoStore.java b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoStore.java
index f97bbbd4e..4626f6e00 100644
--- a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoStore.java
+++ b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoStore.java
@@ -290,6 +290,27 @@ public abstract class OmemoStore<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey, T_
      */
     public abstract void removeOmemoIdentityKey(OmemoDevice userDevice, OmemoDevice contactsDevice);
 
+    /**
+     * Store the number of messages we sent to a device since we last received a message back.
+     * This counter gets reset to 0 whenever we receive a message from the contacts device.
+     *
+     * @param userDevice our omemoDevice.
+     * @param contactsDevice device of which we want to set the message counter.
+     * @param counter counter value.
+     */
+    public abstract void storeOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice, int counter);
+
+    /**
+     * Return the current value of the message counter.
+     * This counter represents the number of message we sent to the contactsDevice without getting a reply back.
+     * The default value for this counter is 0.
+     *
+     * @param userDevice our omemoDevice
+     * @param contactsDevice device of which we want to get the message counter.
+     * @return counter value.
+     */
+    public abstract int loadOmemoMessageCounter(OmemoDevice userDevice, OmemoDevice contactsDevice);
+
     /**
      * Set the date of the last message that was received from a device.
      *
diff --git a/smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoStoreTest.java b/smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoStoreTest.java
index ef5946f82..a4e8bd58e 100644
--- a/smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoStoreTest.java
+++ b/smack-omemo/src/test/java/org/jivesoftware/smackx/omemo/OmemoStoreTest.java
@@ -314,6 +314,13 @@ public abstract class OmemoStoreTest<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey
         assertNull(session);
     }
 
+    @Test
+    public void loadStoreMessageCounterTest() {
+        assertEquals(0, store.loadOmemoMessageCounter(alice, bob));
+        store.storeOmemoMessageCounter(alice, bob, 20);
+        assertEquals(20, store.loadOmemoMessageCounter(alice, bob));
+    }
+
     @Test
     public void getFingerprint() throws IOException, CorruptedOmemoKeyException {
         assertNull("Method must return null for a non-existent fingerprint.", store.getFingerprint(alice));