From 50b41c63aff5e17b3ab5b97e3a00d2603ef64165 Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Tue, 21 May 2019 13:06:46 +1000
Subject: [PATCH] Added storing of last message hash.

---
 Pods                                          |  2 +-
 SignalServiceKit/src/Loki/API/LokiAPI.swift   | 23 ++++++++++++
 .../src/Loki/Crypto/OWSPrimaryStorage+Loki.h  | 28 +++++++++++++-
 .../src/Loki/Crypto/OWSPrimaryStorage+Loki.m  | 37 +++++++++++++++++++
 4 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/Pods b/Pods
index e5b45b28d..c1bf4bf2f 160000
--- a/Pods
+++ b/Pods
@@ -1 +1 @@
-Subproject commit e5b45b28d5e8e409c1acf15c32d3734bb3e8acd7
+Subproject commit c1bf4bf2fb27eaacb36e60b17d30f92f626d4368
diff --git a/SignalServiceKit/src/Loki/API/LokiAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI.swift
index 748fcf45d..2618838e3 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI.swift
@@ -102,3 +102,26 @@ private extension Promise {
         }
     }
 }
+
+// MARK: Last Hash
+
+extension LokiAPI {
+    
+    private var primaryStorage: OWSPrimaryStorage {
+        return OWSPrimaryStorage.shared()
+    }
+    
+    fileprivate func updateLastHash(for node: Target, hash: String, expiresAt: UInt64) {
+        primaryStorage.dbReadWriteConnection.readWrite { transaction in
+            self.primaryStorage.setLastMessageHash(hash, expiresAt: expiresAt, serviceNode: node.address, transaction: transaction)
+        }
+    }
+    
+    fileprivate func getLastHash(for node: Target) -> String? {
+        var lastHash: String? = nil
+        primaryStorage.dbReadWriteConnection.readWrite { transaction in
+            lastHash = self.primaryStorage.getLastMessageHash(forServiceNode: node.address, transaction: transaction)
+        }
+        return lastHash
+    }
+}
diff --git a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.h
index 9a9a694b4..8f97d5e05 100644
--- a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.h
+++ b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.h
@@ -1,8 +1,9 @@
 #import "OWSPrimaryStorage.h"
-#import "PreKeyRecord.h"
-#import "PreKeyBundle.h"
 #import <YapDatabase/YapDatabase.h>
 
+@class PreKeyRecord;
+@class PreKeyBundle;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface OWSPrimaryStorage (Loki)
@@ -72,6 +73,29 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)removePreKeyBundleForContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
+# pragma mark - Last Hash
+
+/**
+ Get the last message hash for the given service node.
+ This function will check the stored last hash and remove it if the `expireAt` has already passed.
+
+ @param serviceNode The service node id
+ @param transaction A read write transaction
+ @return The last hash or nil if it doesn't exist
+ */
+- (NSString *_Nullable)getLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction;
+
+/**
+ Set the last message hash for the given service node.
+ This will override any previous hashes stored for the given service node.
+
+ @param hash The last message hash
+ @param expiresAt The time the message expires on the server
+ @param serviceNode The service node
+ @param transaction A read write transaction
+ */
+- (void)setLastMessageHash:(NSString *)hash expiresAt:(u_int64_t)expiresAt serviceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
index 125679a53..1b678d205 100644
--- a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
+++ b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
@@ -4,6 +4,9 @@
 #import "OWSPrimaryStorage+keyFromIntLong.h"
 #import "OWSDevice.h"
 #import "OWSIdentityManager.h"
+#import "NSDate+OWS.h"
+#import "PreKeyRecord.h"
+#import "PreKeyBundle.h"
 #import "TSAccountManager.h"
 #import "TSPreKeyManager.h"
 #import "YapDatabaseConnection+OWS.h"
@@ -13,6 +16,7 @@
 #define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
 #define LokiPreKeyContactCollection @"LokiPreKeyContactCollection"
 #define LokiPreKeyBundleCollection @"LokiPreKeyBundleCollection"
+#define LokiLastHashCollection @"LokiLastHashCollection"
 
 @implementation OWSPrimaryStorage (Loki)
 
@@ -118,4 +122,37 @@
     [transaction removeObjectForKey:pubKey inCollection:LokiPreKeyBundleCollection];
 }
 
+# pragma mark - Last Hash
+
+- (NSString *_Nullable)getLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    NSDictionary *_Nullable dict = [transaction objectForKey:serviceNode inCollection:LokiLastHashCollection];
+    if (!dict) { return nil; }
+    
+    NSString *_Nullable hash = dict[@"hash"];
+    if (!hash) { return nil; }
+    
+    // Check if the hash isn't expired
+    uint64_t now = NSDate.ows_millisecondTimeStamp;
+    NSNumber *_Nullable expiresAt = dict[@"expiresAt"];
+    if (expiresAt && expiresAt.unsignedLongLongValue <= now) {
+        // The last message has expired from the storage server
+        [self removeLastMessageHashForServiceNode:serviceNode transaction:transaction];
+        return nil;
+    }
+    
+    return hash;
+}
+
+- (void)setLastMessageHash:(NSString *)hash expiresAt:(u_int64_t)expiresAt serviceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    NSDictionary *dict = @{
+                           @"hash": hash,
+                           @"expiresAt": @(expiresAt)
+                           };
+    [transaction setObject:dict forKey:serviceNode inCollection:LokiLastHashCollection];
+}
+
+- (void)removeLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    [transaction removeObjectForKey:serviceNode inCollection:LokiLastHashCollection];
+}
+
 @end