From 247fbfa6b11ccd3cb76c9bfde28c1720737a6ff1 Mon Sep 17 00:00:00 2001
From: Stefan Ceriu <stefanc@matrix.org>
Date: Mon, 30 Aug 2021 16:11:32 +0300
Subject: [PATCH] vector-im/element-ios/issues/4406  - Fixes following code
 review.

---
 MatrixSDK/Crypto/Dehydration/MXDehydrationService.m |  2 +-
 MatrixSDK/Crypto/MXCrypto.m                         |  8 ++++++--
 MatrixSDK/MXRestClient.h                            |  5 ++++-
 MatrixSDKTests/MXRestClientTests.m                  | 13 +------------
 4 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/MatrixSDK/Crypto/Dehydration/MXDehydrationService.m b/MatrixSDK/Crypto/Dehydration/MXDehydrationService.m
index bff08be9bf..1461bbc937 100644
--- a/MatrixSDK/Crypto/Dehydration/MXDehydrationService.m
+++ b/MatrixSDK/Crypto/Dehydration/MXDehydrationService.m
@@ -268,7 +268,7 @@ - (void)uploadDeviceInfo:(MXDeviceInfo*)deviceInfo
     }
 
     MXWeakify(self);
-    [restClient uploadKeys:deviceInfo.JSONDictionary oneTimeKeys:oneTimeJson forDeviceWithId:deviceInfo.deviceId success:^(MXKeysUploadResponse *keysUploadResponse) {
+    [restClient uploadKeys:deviceInfo.JSONDictionary oneTimeKeys:oneTimeJson fallbackKeys:nil forDeviceWithId:deviceInfo.deviceId success:^(MXKeysUploadResponse *keysUploadResponse) {
         [account markOneTimeKeysAsPublished];
         MXLogDebug(@"[MXDehydrationManager] dehydration done succesfully:\n device ID = %@\n ed25519 = %@\n curve25519 = %@", deviceInfo.deviceId, account.identityKeys[@"ed25519"], account.identityKeys[@"curve25519"]);
         MXStrongifyAndReturnIfNil(self);
diff --git a/MatrixSDK/Crypto/MXCrypto.m b/MatrixSDK/Crypto/MXCrypto.m
index 3370a7ba45..45687c2a98 100644
--- a/MatrixSDK/Crypto/MXCrypto.m
+++ b/MatrixSDK/Crypto/MXCrypto.m
@@ -2930,10 +2930,14 @@ - (MXHTTPOperation *)generateAndUploadFallbackKey
     
     MXWeakify(self);
     return [_matrixRestClient uploadKeys:nil oneTimeKeys:nil fallbackKeys:fallbackKeyJson success:^(MXKeysUploadResponse *keysUploadResponse) {
-        weakself.uploadFallbackKeyOperation = nil;
+        MXStrongifyAndReturnIfNil(self);
+        
+        self.uploadFallbackKeyOperation = nil;
         MXLogDebug(@"[MXCrypto] generateAndUploadFallbackKey: Finished uploading fallback key.");
     } failure:^(NSError *error) {
-        weakself.uploadFallbackKeyOperation = nil;
+        MXStrongifyAndReturnIfNil(self);
+        
+        self.uploadFallbackKeyOperation = nil;
         MXLogError(@"[MXCrypto] generateAndUploadFallbackKey: Failed uploading fallback key.");
     }];
 }
diff --git a/MatrixSDK/MXRestClient.h b/MatrixSDK/MXRestClient.h
index 07d58955a8..595d47aa0e 100644
--- a/MatrixSDK/MXRestClient.h
+++ b/MatrixSDK/MXRestClient.h
@@ -2155,6 +2155,7 @@ Get the maximum size a media upload can be in bytes.
 
  @param deviceKeys the device keys to send.
  @param oneTimeKeys the one-time keys to send.
+ @param fallbackKeys the fallback keys to send.
  @param deviceId ID of the device the keys belong to. Nil to upload keys to the device of the current session.
 
  @param success A block object called when the operation succeeds.
@@ -2162,7 +2163,9 @@ Get the maximum size a media upload can be in bytes.
 
  @return a MXHTTPOperation instance.
  */
-- (MXHTTPOperation*)uploadKeys:(NSDictionary*)deviceKeys oneTimeKeys:(NSDictionary*)oneTimeKeys
+- (MXHTTPOperation*)uploadKeys:(NSDictionary*)deviceKeys
+                   oneTimeKeys:(NSDictionary*)oneTimeKeys
+                  fallbackKeys:(NSDictionary *)fallbackKeys
                forDeviceWithId:(NSString*)deviceId
                        success:(void (^)(MXKeysUploadResponse *keysUploadResponse))success
                        failure:(void (^)(NSError *error))failure;
diff --git a/MatrixSDKTests/MXRestClientTests.m b/MatrixSDKTests/MXRestClientTests.m
index 886b3b334d..e891ec13ec 100644
--- a/MatrixSDKTests/MXRestClientTests.m
+++ b/MatrixSDKTests/MXRestClientTests.m
@@ -1581,21 +1581,10 @@ - (void)testOneTimeKeys
         // Upload the device keys
         [bobRestClient uploadKeys:nil oneTimeKeys:otks fallbackKeys:nil success:^(MXKeysUploadResponse *keysUploadResponse) {
             XCTAssert(keysUploadResponse.oneTimeKeyCounts);
-            
             XCTAssertEqual(keysUploadResponse.oneTimeKeyCounts[@"curve25519"].unsignedIntValue, 2, @"Key count must be 2 for 'curve25519'");
             XCTAssertEqual([keysUploadResponse oneTimeKeyCountsForAlgorithm:@"curve25519"], 2, @"Key count must be 2 for 'curve25519'");
-            
-            [keysUploadResponse.oneTimeKeyCounts enumerateKeysAndObjectsUsingBlock:^(NSString *key, NSNumber *obj, BOOL *stop) {
-                if ([key isEqualToString:@"curve25519"]) {
-                    return;
-                }
-                
-                XCTAssertEqual(obj.unsignedIntValue, 0, @"Key count must be 0 for any other algorithm.");
-                XCTAssertEqual([keysUploadResponse oneTimeKeyCountsForAlgorithm:key], 0, @"Key count must be 0 for any other algorithm.");
-            }];
-
+            XCTAssertEqual([keysUploadResponse oneTimeKeyCountsForAlgorithm:@"deded"], 0, @"It must response 0 for any other algo");
             [expectation fulfill];
-
         } failure:^(NSError *error) {
             XCTFail(@"The request should not fail - NSError: %@", error);
             [expectation fulfill];