Correctly handle profile key update from incoming messages.

SignalMessaging/profiles/OWSProfileManager.h

@@ -83,7 +83,7 @@ extern const NSUInteger kOWSProfileManager_MaxAvatarDiameter;
                profileNameEncrypted:(nullable NSData *)profileNameEncrypted
                       avatarUrlPath:(nullable NSString *)avatarUrlPath;
 
-- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName profilePictureURL:(NSString *)profilePictureURL with:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName with:(YapDatabaseReadWriteTransaction *)transaction;
 
 #pragma mark - User Interface
 

SignalMessaging/profiles/OWSProfileManager.m

@@ -340,9 +340,18 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
     }
 }
 
-- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName profilePictureURL:(nullable NSString*)profilePictureURL transaction:(YapDatabaseReadWriteTransaction *)transaction
+- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
-    [self updateLocalProfileName:displayName avatarImage:nil success:^{ } failure:^{ }];
+    [self.localUserProfile updateWithProfileName:displayName transaction:transaction];
+}
+
+- (void)updateUserPofileKeyData:(NSData *)profileKeyData avatarURL:(nullable NSString *)avatarURL transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    OWSAES256Key *profileKey = [OWSAES256Key keyWithData:profileKeyData];
+    if (profileKey != nil) {
+        [self.localUserProfile updateWithProfileKey:profileKey transaction:transaction completion:^{
+            [self.localUserProfile updateWithAvatarUrlPath:avatarURL transaction:transaction];
+        }];
+    }
 }
 
 - (nullable NSString *)profilePictureURL
@@ -409,7 +418,7 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
         // We always want to encrypt a profile with a new profile key
         // This ensures that other users know that our profile picture was updated
-        OWSAES256Key newProfileKey = [OWSAES256Key generateRandomKey];
+        OWSAES256Key *newProfileKey = [OWSAES256Key generateRandomKey];
         NSData *_Nullable encryptedAvatarData;
         if (avatarData) {
             encryptedAvatarData = [self encryptProfileData:avatarData profileKey:newProfileKey];
@@ -942,7 +951,7 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
     }];
 }
 
-- (void)setProfileKeyData:(NSData *)profileKeyData forRecipientId:(NSString *)recipientId
+- (void)setProfileKeyData:(NSData *)profileKeyData forRecipientId:(NSString *)recipientId avatarURL:(nullable NSString *)avatarURL
 {
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
         OWSAES256Key *_Nullable profileKey = [OWSAES256Key keyWithData:profileKeyData];
@@ -950,28 +959,35 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
             OWSFailDebug(@"Failed to make profile key for key data");
             return;
         }
-
+        
-        OWSUserProfile *userProfile =
+        OWSUserProfile *userProfile = [OWSUserProfile getOrBuildUserProfileForRecipientId:recipientId dbConnection:self.dbConnection];
-            [OWSUserProfile getOrBuildUserProfileForRecipientId:recipientId dbConnection:self.dbConnection];
+        
-
         OWSAssertDebug(userProfile);
         if (userProfile.profileKey && [userProfile.profileKey.keyData isEqual:profileKey.keyData]) {
             // Ignore redundant update.
             return;
         }
-
+        
         [userProfile clearWithProfileKey:profileKey
                             dbConnection:self.dbConnection
                               completion:^{
-                                  dispatch_async(dispatch_get_main_queue(), ^{
+            dispatch_async(dispatch_get_main_queue(), ^{
-                                      [self.udManager setUnidentifiedAccessMode:UnidentifiedAccessModeUnknown
+                [self.udManager setUnidentifiedAccessMode:UnidentifiedAccessModeUnknown
-                                                                    recipientId:recipientId];
+                                              recipientId:recipientId];
-                                      [self fetchProfileForRecipientId:recipientId];
+                [userProfile updateWithProfileName:userProfile.profileName avatarUrlPath:avatarURL dbConnection:self.dbConnection completion:^{
-                                  });
+                    [self downloadAvatarForUserProfile:userProfile];
-                              }];
+                }];
+                // [self fetchProfileForRecipientId:recipientId];
+            });
+        }];
     });
 }
 
+- (void)setProfileKeyData:(NSData *)profileKeyData forRecipientId:(NSString *)recipientId
+{
+    [self setProfileKeyData:profileKeyData forRecipientId:recipientId avatarURL:nil];
+}
+
 - (nullable NSData *)profileKeyDataForRecipientId:(NSString *)recipientId
 {
     return [self profileKeyForRecipientId:recipientId].keyData;
@@ -1157,7 +1173,7 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
 
             // If we're updating the profile that corresponds to our local number,
             // update the local profile as well.
-            if (userProfile.address.isLocalAddress) {
+            if ([self.tsAccountManager.localNumber isEqualToString:userProfile.recipientId]) {
                 OWSUserProfile *localUserProfile = self.localUserProfile;
                 OWSAssertDebug(localUserProfile);
 
@@ -1229,17 +1245,10 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
     });
 }
 
-- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName profilePictureURL:(NSString *)profilePictureURL with:(YapDatabaseReadWriteTransaction *)transaction
+- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName with:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSUserProfile *userProfile = [OWSUserProfile getOrBuildUserProfileForRecipientId:contactID transaction:transaction];
-    NSString *oldProfilePictureURL = userProfile.avatarUrlPath;
+    [userProfile updateWithProfileName:displayName avatarUrlPath:userProfile.avatarUrlPath avatarFileName:userProfile.avatarFileName transaction:transaction completion:nil];
-    // Note: we keep using the old file name until we have the new one
-    // (otherwise the profile picuture would disspear for a short time)
-    NSString *oldAvatarFileName = userProfile.avatarFileName;
-    [userProfile updateWithProfileName:displayName avatarUrlPath:profilePictureURL avatarFileName:oldAvatarFileName transaction:transaction completion:nil];
-    if (profilePictureURL && ![oldProfilePictureURL isEqual:profilePictureURL]) {
-        [self downloadAvatarForUserProfile:userProfile];
-    }
 }
 
 - (BOOL)isNullableDataEqual:(NSData *_Nullable)left toData:(NSData *_Nullable)right

SignalMessaging/profiles/OWSUserProfile.h

@@ -73,10 +73,17 @@ extern NSString *const kLocalProfileUniqueId;
                 dbConnection:(YapDatabaseConnection *)dbConnection
                   completion:(nullable OWSUserProfileCompletion)completion;
 
+- (void)updateWithProfileKey:(OWSAES256Key *)profileKey
+                 transaction:(YapDatabaseReadWriteTransaction *)transaction
+                  completion:(nullable OWSUserProfileCompletion)completion;
+
 - (void)clearWithProfileKey:(OWSAES256Key *)profileKey
                dbConnection:(YapDatabaseConnection *)dbConnection
                  completion:(nullable OWSUserProfileCompletion)completion;
 
+- (void)updateWithProfileName:(nullable NSString *)profileName transaction:(YapDatabaseReadWriteTransaction*)transaction;
+- (void)updateWithAvatarUrlPath:(NSString *)avatarUrlPath transaction:(YapDatabaseReadWriteTransaction*)transaction;
+
 #pragma mark - Profile Avatars Directory
 
 + (NSString *)profileAvatarFilepathWithFilename:(NSString *)filename;

SignalMessaging/profiles/OWSUserProfile.m

@@ -316,6 +316,20 @@ NSString *const kLocalProfileUniqueId = @"kLocalProfileUniqueId";
             completion:completion];
 }
 
+- (void)updateWithProfileName:(nullable NSString *)profileName transaction:(YapDatabaseReadWriteTransaction*)transaction
+{
+    [self updateWithProfileName:profileName avatarUrlPath:self.avatarUrlPath avatarFileName:self.avatarFileName transaction:transaction completion:nil];
+}
+
+- (void)updateWithAvatarUrlPath:(NSString *)avatarUrlPath transaction:(YapDatabaseReadWriteTransaction*)transaction {
+    [self applyChanges:^(OWSUserProfile *userProfile) {
+        [userProfile setAvatarUrlPath:avatarUrlPath];
+    }
+           functionName:__PRETTY_FUNCTION__
+           transaction:transaction
+            completion:nil];
+}
+
 - (void)updateWithAvatarUrlPath:(nullable NSString *)avatarUrlPath
                  avatarFileName:(nullable NSString *)avatarFileName
                    dbConnection:(YapDatabaseConnection *)dbConnection

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -554,17 +554,9 @@ NS_ASSUME_NONNULL_BEGIN
             return;
         }
     }
+    
+    [self handleProfileKeyUpdateIfNeeded:dataMessage recipientId:envelope.source];
 
-    if ([dataMessage hasProfileKey]) {
-        NSData *profileKey = [dataMessage profileKey];
-        NSString *recipientId = envelope.source;
-        if (profileKey.length == kAES256_KeyByteLength) {
-            [self.profileManager setProfileKeyData:profileKey forRecipientId:recipientId];
-        } else {
-            OWSFailDebug(
-                @"Unexpected profile key length:%lu on message from:%@", (unsigned long)profileKey.length, recipientId);
-        }
-    }
 
     if (dataMessage.group) {
         TSGroupThread *_Nullable groupThread =
@@ -932,10 +924,14 @@ NS_ASSUME_NONNULL_BEGIN
         // Loki: Try to update using the provided profile
         if (wasSentByMasterDevice) {
             dispatch_async(dispatch_get_main_queue(), ^{
-                SSKProtoDataMessageLokiProfile *profile = syncMessage.sent.message.profile;
+                SSKProtoDataMessage *dataMessage = syncMessage.sent.message;
+                SSKProtoDataMessageLokiProfile *profile = dataMessage.profile;
                 NSString *displayName = profile.displayName;
                 NSString *profilePictureURL = profile.profilePicture;
-                [self.profileManager updateUserProfileWithDisplayName:displayName profilePictureURL:profilePictureURL transaction:transaction];
+                [self.profileManager updateUserProfileWithDisplayName:displayName transaction:transaction];
+                if ([dataMessage hasProfileKey]) {
+                    [self.profileManager updateUserPofileKeyData:dataMessage.profileKey avatarURL:profilePictureURL transaction:transaction];
+                }
             });
         }
 
@@ -1203,9 +1199,14 @@ NS_ASSUME_NONNULL_BEGIN
             envelopeAddress(envelope));
         return;
     }
+    
+    if (dataMessage.profile == nil) {
+        OWSFailDebug(@"received profile key message without loki profile attached from: %@", envelopeAddress(envelope));
+        return;
+    }
 
     id<ProfileManagerProtocol> profileManager = SSKEnvironment.shared.profileManager;
-    [profileManager setProfileKeyData:profileKey forRecipientId:recipientId];
+    [profileManager setProfileKeyData:profileKey forRecipientId:recipientId avatarURL:dataMessage.profile.profilePicture];
 }
 
 - (void)handleUnlinkDeviceMessageWithEnvelope:(SSKProtoEnvelope *)envelope dataMessage:(SSKProtoDataMessage *)dataMessage transaction:(YapDatabaseReadWriteTransaction *)transaction
@@ -1388,9 +1389,9 @@ NS_ASSUME_NONNULL_BEGIN
         TSContactThread *thread =
             [TSContactThread getOrCreateThreadWithContactId:hexEncodedPublicKey transaction:transaction];
 
-        NSString *profilePictureURL = dataMessage.profile.profilePicture;
+        // Only set the display name here, the logic for updating profile pictures is handled when we're setting profile key
         NSString *displayName = dataMessage.profile.displayName;
-        [self.profileManager updateProfileForContactWithID:thread.contactIdentifier displayName:displayName profilePictureURL:profilePictureURL with:transaction];
+        [self.profileManager updateProfileForContactWithID:thread.contactIdentifier displayName:displayName with:transaction];
 
         switch (dataMessage.group.type) {
             case SSKProtoGroupContextTypeUpdate: {
@@ -1588,13 +1589,8 @@ NS_ASSUME_NONNULL_BEGIN
         if (rawDisplayName != nil && rawDisplayName.length > 0) {
             displayName = [NSString stringWithFormat:@"%@ (...%@)", rawDisplayName, [incomingMessage.authorId substringFromIndex:incomingMessage.authorId.length - 8]];
         }
-        NSString *rawProfilePictureURL = dataMessage.profile.profilePicture;
+        [self.profileManager updateProfileForContactWithID:thread.contactIdentifier displayName:displayName with:transaction];
-        NSString *profilePictureURL = nil;
+        [self handleProfileKeyUpdateIfNeeded:dataMessage recipientId:thread.contactIdentifier];
-        if (rawProfilePictureURL != nil && rawProfilePictureURL.length > 0) {
-            profilePictureURL = rawProfilePictureURL;
-        }
-
-        [self.profileManager updateProfileForContactWithID:thread.contactIdentifier displayName:displayName profilePictureURL:profilePictureURL with:transaction];
 
         // Loki: Parse Loki specific properties if needed
         if (envelope.isPtpMessage) { incomingMessage.isP2P = YES; }
@@ -1635,6 +1631,19 @@ NS_ASSUME_NONNULL_BEGIN
     }
 }
 
+- (void)handleProfileKeyUpdateIfNeeded:(SSKProtoDataMessage *)dataMessage recipientId:(NSString *)recipientId {
+    if ([dataMessage hasProfileKey]) {
+        NSData *profileKey = [dataMessage profileKey];
+        NSString *url = dataMessage.profile != nil ? dataMessage.profile.profilePicture : nil;
+        if (profileKey.length == kAES256_KeyByteLength) {
+            [self.profileManager setProfileKeyData:profileKey forRecipientId:recipientId avatarURL:url];
+        } else {
+            OWSFailDebug(
+                         @"Unexpected profile key length:%lu on message from:%@", (unsigned long)profileKey.length, recipientId);
+        }
+    }
+}
+
 - (BOOL)canFriendRequestBeAutoAcceptedForThread:(TSContactThread *)thread transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     NSString *senderHexEncodedPublicKey = thread.contactIdentifier;

SignalServiceKit/src/Protocols/ProfileManagerProtocol.h

@@ -5,6 +5,7 @@
 @class OWSAES256Key;
 @class TSThread;
 @class YapDatabaseReadWriteTransaction;
+@class OWSUserProfile;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -17,8 +18,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (nullable NSString *)profilePictureURL;
 
 - (nullable NSData *)profileKeyDataForRecipientId:(NSString *)recipientId;
-- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName profilePictureURL:(nullable NSString*)profilePictureURL transaction:(YapDatabaseReadWriteTransaction *)transaction;
 - (void)setProfileKeyData:(NSData *)profileKeyData forRecipientId:(NSString *)recipientId;
+- (void)setProfileKeyData:(NSData *)profileKeyData forRecipientId:(NSString *)recipientId avatarURL:(nullable NSString *)avatarURL;
 
 - (BOOL)isUserInProfileWhitelist:(NSString *)recipientId;
 
@@ -31,7 +32,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)fetchProfileForRecipientId:(NSString *)recipientId;
 
-- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName profilePictureURL:(NSString *)profilePictureURL with:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName transaction:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)updateUserPofileKeyData:(NSData *)profileKeyData avatarURL:(nullable NSString *)avatarURL transaction:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)updateProfileForContactWithID:(NSString *)contactID displayName:(NSString *)displayName with:(YapDatabaseReadWriteTransaction *)transaction;
 
 @end
 

SignalServiceKit/src/TestUtils/OWSFakeProfileManager.m

@@ -48,7 +48,7 @@ NS_ASSUME_NONNULL_BEGIN
     return _localProfileKey;
 }
 
-- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName profilePictureURL:(nullable NSString*)profilePictureURL transaction:(YapDatabaseReadWriteTransaction *)transaction
+- (void)updateUserProfileWithDisplayName:(nullable NSString*)displayName transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     // Do nothing
 }