Fix | Share room keys with dehydrated devices with rust stack

MatrixSDK/Categories/MXKeysQueryResponse+Extensions.swift

@@ -39,3 +39,27 @@ extension MXKeysQueryResponse : MXSummable {
         return keysQueryResponse as! Self
     }
 }
+
+
+extension MXKeysQueryResponseRaw : MXSummable {
+
+    public static func +(lhs: MXKeysQueryResponseRaw, rhs: MXKeysQueryResponseRaw) -> Self {
+        let keysQueryResponse = MXKeysQueryResponseRaw()
+
+        // Casts to original objc NSDictionary are annoying
+        // but we want to reuse our implementation of NSDictionary.+
+        let deviceKeysMap = (lhs.deviceKeys as NSDictionary? ?? NSDictionary())
+        + (rhs.deviceKeys as NSDictionary? ?? NSDictionary())
+        keysQueryResponse.deviceKeys = deviceKeysMap as? [String : Any]
+
+        let crossSigningKeys = (lhs.crossSigningKeys as NSDictionary? ?? NSDictionary())
+            + (rhs.crossSigningKeys as NSDictionary? ?? NSDictionary())
+        keysQueryResponse.crossSigningKeys = crossSigningKeys as? [String: MXCrossSigningInfo]
+
+        let failures = (lhs.failures as NSDictionary? ?? NSDictionary())
+            + (rhs.failures as NSDictionary? ?? NSDictionary())
+        keysQueryResponse.failures = failures as? [AnyHashable : Any]
+
+        return keysQueryResponse as! Self
+    }
+}

MatrixSDK/Categories/MXRestClient+Extensions.swift

@@ -89,4 +89,75 @@ public extension MXRestClient {
 
         return operation
     }
+
+    /// Download users keys by chunks.
+    ///
+    /// - Parameters:
+    ///   - users: list of users to get keys for.
+    ///   - token: sync token to pass in the query request, to help.
+    ///   - chunkSize: max number of users to ask for in one CS API request.
+    ///   - success: A block object called when the operation succeeds.
+    ///   - failure: A block object called when the operation fails.
+    /// - Returns: a MXHTTPOperation instance.
+    func downloadKeysByChunkRaw(forUsers users: [String],
+                             token: String?,
+                             chunkSize: Int = 250,
+                             success: @escaping (_ keysQueryResponse: MXKeysQueryResponseRaw) -> Void,
+                             failure: @escaping (_ error: NSError?) -> Void) -> MXHTTPOperation {
+
+        // Do not chunk if not needed
+        if users.count <= chunkSize {
+            return self.downloadKeysRaw(forUsers: users, token: token) { response in
+                switch response {
+                    case .success(let keysQueryResponse):
+                        success(keysQueryResponse)
+                    case .failure(let error):
+                        failure(error as NSError)
+                }
+            }
+        }
+
+        MXLog.debug("[MXRestClient+Extensions] downloadKeysByChunk: \(users.count) users with chunkSize:\(chunkSize)")
+
+        // An arbitrary MXHTTPOperation. It will not cancel requests
+        // but it will avoid to call callbacks in case of a cancellation is requested
+        let operation = MXHTTPOperation()
+
+        let group = DispatchGroup()
+        var responses = [MXResponse<MXKeysQueryResponseRaw>]()
+        users.chunked(into: chunkSize).forEach { chunkedUsers in
+            group.enter()
+            self.downloadKeysRaw(forUsers: chunkedUsers, token: token) { response in
+                switch response {
+                    case .success(let keysQueryResponse):
+                        MXLog.debug("[MXRestClient+Extensions] downloadKeysByChunk: Got intermediate response. Got device keys for %@ users. Got cross-signing keys for %@ users \(String(describing: keysQueryResponse.deviceKeys.keys.count)) \(String(describing: keysQueryResponse.crossSigningKeys.count))")
+                    case .failure(let error):
+                        MXLog.debug("[MXRestClient+Extensions] downloadKeysByChunk: Got intermediate error. Error: \(error)")
+                }
+
+                responses.append(response)
+                group.leave()
+            }
+        }
+
+        group.notify(queue: self.completionQueue) {
+            MXLog.debug("[MXRestClient+Extensions] downloadKeysByChunk: Got all responses")
+
+            guard operation.isCancelled == false else {
+                MXLog.debug("[MXRestClient+Extensions] downloadKeysByChunk: Request was cancelled")
+                return
+            }
+
+            // Gather all responses in one
+            let response = responses.reduce(.success(MXKeysQueryResponseRaw()), +)
+            switch response {
+                case .success(let keysQueryResponse):
+                    success(keysQueryResponse)
+                case .failure(let error):
+                    failure(error as NSError)
+            }
+        }
+
+        return operation
+    }
 }

MatrixSDK/Contrib/Swift/MXRestClient.swift

@@ -1848,6 +1848,10 @@ public extension MXRestClient {
         return __downloadKeys(forUsers: userIds, token: token, success: currySuccess(completion), failure: curryFailure(completion))
     }
 
+    @nonobjc @discardableResult func downloadKeysRaw(forUsers userIds: [String], token: String? = nil, completion: @escaping (_ response: MXResponse<MXKeysQueryResponseRaw>) -> Void) -> MXHTTPOperation {
+        return __downloadKeysRaw(forUsers: userIds, token: token, success: currySuccess(completion), failure: curryFailure(completion))
+    }
+
 
     /**
      Claim one-time keys.

MatrixSDK/Crypto/CryptoMachine/MXCryptoRequests.swift

@@ -21,13 +21,13 @@ import MatrixSDKCrypto
 /// to the native REST API client
 struct MXCryptoRequests {
     private let restClient: MXRestClient
-    private let queryScheduler: MXKeysQueryScheduler<MXKeysQueryResponse>
+    private let queryScheduler: MXKeysQueryScheduler<MXKeysQueryResponseRaw>
 
     init(restClient: MXRestClient) {
         self.restClient = restClient
         self.queryScheduler = .init { users in
             try await performCallbackRequest { completion in
-                _ = restClient.downloadKeysByChunk(
+                _ = restClient.downloadKeysByChunkRaw(
                     forUsers: users,
                     token: nil,
                     success: {
@@ -96,7 +96,7 @@ struct MXCryptoRequests {
         }
     }
 
-    func queryKeys(users: [String]) async throws -> MXKeysQueryResponse {
+    func queryKeys(users: [String]) async throws -> MXKeysQueryResponseRaw {
         try await queryScheduler.query(users: Set(users))
     }
 

MatrixSDK/JSONModels/MXJSONModels.h

@@ -1125,6 +1125,25 @@ FOUNDATION_EXPORT NSString *const kMXPushRuleScopeStringGlobal;
 
 @end
 
+@interface MXKeysQueryResponseRaw : MXJSONModel
+
+    /**
+     The device keys per devices per users.
+     */
+    @property (nonatomic) NSDictionary<NSString *, id> *deviceKeys;
+
+    /**
+     Cross-signing keys per users.
+     */
+    @property (nonatomic) NSDictionary<NSString*, MXCrossSigningInfo*> *crossSigningKeys;
+
+    /**
+     The failures sorted by homeservers.
+    */
+    @property (nonatomic) NSDictionary *failures;
+
+@end
+
 /**
  `MXKeysClaimResponse` represents the response to /keys/claim request made by
  [MXRestClient claimOneTimeKeysForUsersDevices].

MatrixSDK/JSONModels/MXJSONModels.m

@@ -1247,6 +1247,101 @@ - (NSDictionary *)JSONDictionary
 
 @end
 
+@interface MXKeysQueryResponseRaw ()
+@end
+
+@implementation MXKeysQueryResponseRaw
+
++ (id)modelFromJSON:(NSDictionary *)JSONDictionary
+{
+    MXKeysQueryResponseRaw *keysQueryResponse = [[MXKeysQueryResponseRaw alloc] init];
+    if (keysQueryResponse)
+    {
+
+        if ([JSONDictionary[@"device_keys"] isKindOfClass:NSDictionary.class])
+        {
+            keysQueryResponse.deviceKeys = JSONDictionary[@"device_keys"];
+        }
+
+        MXJSONModelSetDictionary(keysQueryResponse.failures, JSONDictionary[@"failures"]);
+
+        // Extract cross-signing keys
+        NSMutableDictionary *crossSigningKeys = [NSMutableDictionary dictionary];
+
+        // Gather all of them by type by user
+        NSDictionary<NSString*, NSDictionary<NSString*, MXCrossSigningKey*>*> *allKeys =
+        @{
+          MXCrossSigningKeyType.master: [self extractUserKeysFromJSON:JSONDictionary[@"master_keys"]] ?: @{},
+          MXCrossSigningKeyType.selfSigning: [self extractUserKeysFromJSON:JSONDictionary[@"self_signing_keys"]] ?: @{},
+          MXCrossSigningKeyType.userSigning: [self extractUserKeysFromJSON:JSONDictionary[@"user_signing_keys"]] ?: @{},
+          };
+
+        // Package them into a `userId -> MXCrossSigningInfo` dictionary
+        for (NSString *keyType in allKeys)
+        {
+            NSDictionary<NSString*, MXCrossSigningKey*> *keys = allKeys[keyType];
+            for (NSString *userId in keys)
+            {
+                MXCrossSigningInfo *crossSigningInfo = crossSigningKeys[userId];
+                if (!crossSigningInfo)
+                {
+                    crossSigningInfo = [[MXCrossSigningInfo alloc] initWithUserId:userId];
+                    crossSigningKeys[userId] = crossSigningInfo;
+                }
+
+                [crossSigningInfo addCrossSigningKey:keys[userId] type:keyType];
+            }
+        }
+
+        keysQueryResponse.crossSigningKeys = crossSigningKeys;
+    }
+
+    return keysQueryResponse;
+}
+
++ (NSDictionary<NSString*, MXCrossSigningKey*>*)extractUserKeysFromJSON:(NSDictionary *)keysJSONDictionary
+{
+    NSMutableDictionary<NSString*, MXCrossSigningKey*> *keys = [NSMutableDictionary dictionary];
+    for (NSString *userId in keysJSONDictionary)
+    {
+        MXCrossSigningKey *key;
+        MXJSONModelSetMXJSONModel(key, MXCrossSigningKey, keysJSONDictionary[userId]);
+        if (key)
+        {
+            keys[userId] = key;
+        }
+    }
+
+    if (!keys.count)
+    {
+        keys = nil;
+    }
+
+    return keys;
+}
+
+- (NSDictionary *)JSONDictionary
+{
+
+    NSMutableDictionary *master = [[NSMutableDictionary alloc] init];
+    NSMutableDictionary *selfSigning = [[NSMutableDictionary alloc] init];
+    NSMutableDictionary *userSigning = [[NSMutableDictionary alloc] init];
+    for (NSString *userId in self.crossSigningKeys) {
+        master[userId] = self.crossSigningKeys[userId].masterKeys.JSONDictionary.copy;
+        selfSigning[userId] = self.crossSigningKeys[userId].selfSignedKeys.JSONDictionary.copy;
+        userSigning[userId] = self.crossSigningKeys[userId].userSignedKeys.JSONDictionary.copy;
+    }
+
+    return @{
+        @"device_keys": self.deviceKeys.copy ?: @{},
+        @"failures": self.failures.copy ?: @{},
+        @"master_keys": master.copy ?: @{},
+        @"self_signing_keys": selfSigning.copy ?: @{},
+        @"user_signing_keys": userSigning.copy ?: @{}
+    };
+}
+
+@end
 @interface MXKeysClaimResponse ()
 
 /**

MatrixSDK/MXRestClient.h

@@ -2485,6 +2485,10 @@ Note: Clients should consider avoiding this endpoint for URLs posted in encrypte
                                  success:(void (^)(MXKeysQueryResponse *keysQueryResponse))success
                                  failure:(void (^)(NSError *error))failure NS_REFINED_FOR_SWIFT;
 
+- (MXHTTPOperation*)downloadKeysRawForUsers:(NSArray<NSString*>*)userIds
+                                   token:(NSString*)token
+                                 success:(void (^)(MXKeysQueryResponseRaw *keysQueryResponse))success
+                                 failure:(void (^)(NSError *error))failure NS_REFINED_FOR_SWIFT;
 /**
  * Claim one-time keys.
 

MatrixSDK/MXRestClient.m

@@ -4925,6 +4925,51 @@ - (MXHTTPOperation*)downloadKeysForUsers:(NSArray<NSString*>*)userIds
                                  }];
 }
 
+- (MXHTTPOperation*)downloadKeysRawForUsers:(NSArray<NSString*>*)userIds
+                                   token:(NSString *)token
+                                 success:(void (^)(MXKeysQueryResponseRaw *keysQueryResponse))success
+                                 failure:(void (^)(NSError *error))failure
+{
+    NSString *path = [NSString stringWithFormat:@"%@/keys/query", kMXAPIPrefixPathR0];
+
+    NSMutableDictionary *downloadQuery = [NSMutableDictionary dictionary];
+    for (NSString *userID in userIds)
+    {
+        downloadQuery[userID] = @[];
+    }
+
+    NSMutableDictionary *parameters = [NSMutableDictionary dictionaryWithDictionary:@{
+                                                                                      @"device_keys": downloadQuery
+                                                                                      }];
+
+    if (token)
+    {
+        parameters[@"token"] = token;
+    }
+
+    MXWeakify(self);
+    return [httpClient requestWithMethod:@"POST"
+                                    path: path
+                              parameters:parameters
+                                 success:^(NSDictionary *JSONResponse) {
+                                     MXStrongifyAndReturnIfNil(self);
+
+                                     if (success)
+                                     {
+                                         __block MXKeysQueryResponseRaw *keysQueryResponse;
+                                         [self dispatchProcessing:^{
+                                             MXJSONModelSetMXJSONModel(keysQueryResponse, MXKeysQueryResponseRaw, JSONResponse);
+                                         } andCompletion:^{
+                                             success(keysQueryResponse);
+                                         }];
+                                     }
+                                 }
+                                 failure:^(NSError *error) {
+                                     MXStrongifyAndReturnIfNil(self);
+                                     [self dispatchFailure:error inBlock:failure];
+                                 }];
+}
+
 - (MXHTTPOperation *)claimOneTimeKeysForUsersDevices:(MXUsersDevicesMap<NSString *> *)usersDevicesKeyTypesMap success:(void (^)(MXKeysClaimResponse *))success failure:(void (^)(NSError *))failure
 {
     NSString *path = [NSString stringWithFormat:@"%@/keys/claim", kMXAPIPrefixPathR0];