From 4a6a02c009acbdcb6d7172ee7009034f77b9e03c Mon Sep 17 00:00:00 2001
From: Michael Kirk <michael.code@endoftheworl.de>
Date: Wed, 24 May 2017 17:30:30 -0700
Subject: [PATCH] Ensure updates don't clobber

// FREEBIE
---
 src/Security/OWSRecipientIdentity.h           | 19 ++++++---
 src/Security/OWSRecipientIdentity.m           | 41 ++++++++++++++++++-
 .../TSStorageManager+IdentityKeyStore.m       |  5 +--
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/src/Security/OWSRecipientIdentity.h b/src/Security/OWSRecipientIdentity.h
index 522257f30..d21699d65 100644
--- a/src/Security/OWSRecipientIdentity.h
+++ b/src/Security/OWSRecipientIdentity.h
@@ -11,10 +11,21 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, readonly) NSString *recipientId;
 @property (nonatomic, readonly) NSData *identityKey;
 @property (nonatomic, readonly) NSDate *createdAt;
-@property (atomic, readonly) BOOL wasSeen;
 @property (nonatomic, readonly) BOOL isFirstKnownKey;
-@property (atomic) BOOL approvedForBlockingUse;
-@property (atomic) BOOL approvedForNonBlockingUse;
+
+#pragma mark - get/set Seen
+
+@property (atomic, readonly) BOOL wasSeen;
+- (void)updateAsSeen;
+
+#pragma mark - get/set Approval
+
+@property (atomic, readonly) BOOL approvedForBlockingUse;
+@property (atomic, readonly) BOOL approvedForNonBlockingUse;
+- (void)updateWithApprovedForBlockingUse:(BOOL)approvedForBlockingUse
+               approvedForNonBlockingUse:(BOOL)approvedForNonBlockingUse;
+
+#pragma mark - Initializers
 
 - (instancetype)initWithUniqueId:(NSString *)uniqueId NS_UNAVAILABLE;
 
@@ -25,8 +36,6 @@ NS_ASSUME_NONNULL_BEGIN
              approvedForBlockingUse:(BOOL)approvedForBlockingUse
           approvedForNonBlockingUse:(BOOL)approvedForNonBlockingUse NS_DESIGNATED_INITIALIZER;
 
-- (void)markAsSeen;
-
 #pragma mark - debug
 
 + (void)printAllIdentities;
diff --git a/src/Security/OWSRecipientIdentity.m b/src/Security/OWSRecipientIdentity.m
index ad7d1e580..23ff04032 100644
--- a/src/Security/OWSRecipientIdentity.m
+++ b/src/Security/OWSRecipientIdentity.m
@@ -6,6 +6,15 @@
 #import "TSStorageManager.h"
 
 NS_ASSUME_NONNULL_BEGIN
+
+@interface OWSRecipientIdentity ()
+
+@property (atomic) BOOL wasSeen;
+@property (atomic) BOOL approvedForBlockingUse;
+@property (atomic) BOOL approvedForNonBlockingUse;
+
+@end
+
 /**
  * Record for a recipients identity key and some meta data around it used to make trust decisions.
  *
@@ -37,9 +46,37 @@ NS_ASSUME_NONNULL_BEGIN
     return self;
 }
 
-- (void)markAsSeen
+- (void)updateAsSeen
+{
+    [self updateWithChangeBlock:^(OWSRecipientIdentity *_Nonnull obj) {
+        obj.wasSeen = YES;
+    }];
+}
+
+- (void)updateWithApprovedForBlockingUse:(BOOL)approvedForBlockingUse
+               approvedForNonBlockingUse:(BOOL)approvedForNonBlockingUse
 {
-    _wasSeen = YES;
+    // Ensure changes are persisted without clobbering any work done on another thread or instance.
+    [self updateWithChangeBlock:^(OWSRecipientIdentity *_Nonnull obj) {
+        obj.approvedForBlockingUse = approvedForBlockingUse;
+        obj.approvedForNonBlockingUse = approvedForNonBlockingUse;
+    }];
+}
+
+- (void)updateWithChangeBlock:(void (^)(OWSRecipientIdentity *obj))changeBlock
+{
+    changeBlock(self);
+
+    [[self class].dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+        OWSRecipientIdentity *latest = [[self class] fetchObjectWithUniqueID:self.uniqueId transaction:transaction];
+        if (latest == nil) {
+            [self saveWithTransaction:transaction];
+            return;
+        }
+
+        changeBlock(latest);
+        [latest saveWithTransaction:transaction];
+    }];
 }
 
 /**
diff --git a/src/Storage/AxolotlStore/TSStorageManager+IdentityKeyStore.m b/src/Storage/AxolotlStore/TSStorageManager+IdentityKeyStore.m
index ad9b5ca4f..bca39d526 100644
--- a/src/Storage/AxolotlStore/TSStorageManager+IdentityKeyStore.m
+++ b/src/Storage/AxolotlStore/TSStorageManager+IdentityKeyStore.m
@@ -125,9 +125,8 @@ const NSTimeInterval kIdentityKeyStoreNonBlockingSecondsThreshold = 5.0;
         }
         
         if ([self isBlockingApprovalRequiredForIdentity:existingIdentity] || [self isNonBlockingApprovalRequiredForIdentity:existingIdentity]) {
-            existingIdentity.approvedForBlockingUse = approvedForBlockingUse;
-            existingIdentity.approvedForNonBlockingUse = approvedForNonBlockingUse;
-            [existingIdentity save];
+            [existingIdentity updateWithApprovedForBlockingUse:approvedForBlockingUse
+                                     approvedForNonBlockingUse:approvedForNonBlockingUse];
             return NO;
         }