keychain/ckks/tests/CloudKitMockXCTest.m

   1 /*
   2  * Copyright (c) 2016 Apple Inc. All Rights Reserved.
   3  *
   4  * @APPLE_LICENSE_HEADER_START@
   5  *
   6  * This file contains Original Code and/or Modifications of Original Code
   7  * as defined in and that are subject to the Apple Public Source License
   8  * Version 2.0 (the 'License'). You may not use this file except in
   9  * compliance with the License. Please obtain a copy of the License at
  10  * http://www.opensource.apple.com/apsl/ and read it before using this
  11  * file.
  12  *
  13  * The Original Code and all software distributed under the License are
  14  * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
  15  * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
  16  * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
  17  * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
  18  * Please see the License for the specific language governing rights and
  19  * limitations under the License.
  20  *
  21  * @APPLE_LICENSE_HEADER_END@
  22  */
  23
  24 #if OCTAGON
  25
  26 #import "CloudKitMockXCTest.h"
  27
  28 #import <ApplePushService/ApplePushService.h>
  29 #import <Foundation/Foundation.h>
  30 #import <CloudKit/CloudKit.h>
  31 #import <CloudKit/CloudKit_Private.h>
  32 #import <CloudKit/CKContainer_Private.h>
  33 #import <OCMock/OCMock.h>
  34
  35 #include "keychain/securityd/Regressions/SecdTestKeychainUtilities.h"
  36 #include <utilities/SecFileLocations.h>
  37 #include "keychain/securityd/SecItemServer.h"
  38
  39 #if NO_SERVER
  40 #include "keychain/securityd/spi.h"
  41 #endif
  42
  43 #include <Security/SecureObjectSync/SOSViews.h>
  44
  45 #include <utilities/SecDb.h>
  46 #include "keychain/securityd/SecItemServer.h"
  47 #include <keychain/ckks/CKKS.h>
  48 #include <keychain/ckks/CKKSViewManager.h>
  49 #include <keychain/ckks/CKKSKeychainView.h>
  50 #include <keychain/ckks/CKKSItem.h>
  51 #include <keychain/ckks/CKKSOutgoingQueueEntry.h>
  52 #include <keychain/ckks/CKKSKey.h>
  53 #include "keychain/ckks/CKKSGroupOperation.h"
  54 #include "keychain/ckks/CKKSLockStateTracker.h"
  55 #include "keychain/ckks/CKKSReachabilityTracker.h"
  56
  57 #import "tests/secdmockaks/mockaks.h"
  58 #import "utilities/SecTapToRadar.h"
  59
  60 #import "MockCloudKit.h"
  61
  62 @interface BoolHolder : NSObject
  63 @property bool state;
  64 @end
  65
  66 @implementation BoolHolder
  67 @end
  68
  69 // Inform OCMock about the internals of CKContainer
  70 @interface CKContainer ()
  71 - (void)_checkSelfCloudServicesEntitlement;
  72 @end
  73
  74
  75 @implementation CloudKitMockXCTest
  76 @synthesize aksLockState = _aksLockState;
  77
  78 + (void)setUp {
  79     // Turn on testing
  80     SecCKKSEnable();
  81     SecCKKSTestsEnable();
  82     SecCKKSSetReduceRateLimiting(true);
  83     [super setUp];
  84
  85 #if NO_SERVER
  86     securityd_init_local_spi();
  87 #endif
  88 }
  89
  90 - (BOOL)isRateLimited:(SecTapToRadar *)ttrRequest
  91 {
  92     return self.isTTRRatelimited;
  93 }
  94
  95 - (BOOL)askUserIfTTR:(SecTapToRadar *)ttrRequest
  96 {
  97     return YES;
  98 }
  99
 100 - (void)triggerTapToRadar:(SecTapToRadar *)ttrRequest
 101 {
 102     [self.ttrExpectation fulfill];
 103 }
 104
 105 - (void)setUp {
 106     [super setUp];
 107
 108     NSString* testName = [self.name componentsSeparatedByString:@" "][1];
 109     testName = [testName stringByReplacingOccurrencesOfString:@"]" withString:@""];
 110     secnotice("ckkstest", "Beginning test %@", testName);
 111
 112     // All tests start with the same flag set.
 113     SecCKKSTestResetFlags();
 114     SecCKKSTestSetDisableSOS(true);
 115
 116     self.silentFetchesAllowed = true;
 117     self.silentZoneDeletesAllowed = false; // Set to true if you want to do any deletes
 118
 119     __weak __typeof(self) weakSelf = self;
 120     self.operationQueue = [[NSOperationQueue alloc] init];
 121     self.operationQueue.maxConcurrentOperationCount = 1;
 122
 123     self.zones = [[NSMutableDictionary alloc] init];
 124
 125     self.apsEnvironment = @"fake APS push string";
 126
 127     // Static variables are a scourge. Let's reset this one...
 128     [OctagonAPSReceiver resetGlobalEnviornmentMap];
 129
 130     self.mockDatabaseExceptionCatcher = OCMStrictClassMock([CKDatabase class]);
 131     self.mockDatabase = OCMStrictClassMock([CKDatabase class]);
 132     self.mockContainerExpectations = OCMStrictClassMock([CKContainer class]);
 133     self.mockContainer = OCMClassMock([CKContainer class]);
 134     OCMStub([self.mockContainer containerWithIdentifier:[OCMArg isKindOfClass:[NSString class]]]).andReturn(self.mockContainer);
 135     OCMStub([self.mockContainer defaultContainer]).andReturn(self.mockContainer);
 136     OCMStub([self.mockContainer alloc]).andReturn(self.mockContainer);
 137     OCMStub([self.mockContainer containerIdentifier]).andReturn(SecCKKSContainerName);
 138     OCMStub([self.mockContainer initWithContainerID: [OCMArg any] options: [OCMArg any]]).andReturn(self.mockContainer);
 139     OCMStub([self.mockContainer privateCloudDatabase]).andReturn(self.mockDatabaseExceptionCatcher);
 140     OCMStub([self.mockContainer serverPreferredPushEnvironmentWithCompletionHandler: ([OCMArg invokeBlockWithArgs:self.apsEnvironment, [NSNull null], nil])]);
 141     OCMStub([self.mockContainer submitEventMetric:[OCMArg any]]).andCall(self, @selector(ckcontainerSubmitEventMetric:));
 142
 143     // Use two layers of mockDatabase here, so we can both add Expectations and catch the exception (instead of crash) when one fails.
 144     OCMStub([self.mockDatabaseExceptionCatcher addOperation:[OCMArg any]]).andCall(self, @selector(ckdatabaseAddOperation:));
 145
 146     // If you want to change this, you'll need to update the mock
 147     _ckDeviceID = [NSString stringWithFormat:@"fake-cloudkit-device-id-%@", testName];
 148     OCMStub([self.mockContainer fetchCurrentDeviceIDWithCompletionHandler: ([OCMArg invokeBlockWithArgs:self.ckDeviceID, [NSNull null], nil])]);
 149
 150     self.accountStatus = CKAccountStatusAvailable;
 151     self.iCloudHasValidCredentials = YES;
 152
 153     self.fakeHSA2AccountStatus = CKKSAccountStatusAvailable;
 154
 155     // Inject a fake operation dependency so we won't respond with the CloudKit account status immediately
 156     // The CKKSAccountStateTracker won't send any login/logout calls without that information, so this blocks all CKKS setup
 157     self.ckaccountHoldOperation = [NSBlockOperation named:@"ckaccount-hold" withBlock:^{
 158         secnotice("ckks", "CKKS CK account status test hold released");
 159     }];
 160
 161     OCMStub([self.mockContainer accountStatusWithCompletionHandler:
 162                     [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (CKAccountStatus accountStatus,
 163                                                                       NSError * _Nullable error)) {
 164
 165         if(passedBlock) {
 166             __strong __typeof(self) strongSelf = weakSelf;
 167             NSBlockOperation* fulfillBlock = [NSBlockOperation named:@"account-status-completion" withBlock: ^{
 168                 passedBlock(weakSelf.accountStatus, nil);
 169             }];
 170             [fulfillBlock addDependency: strongSelf.ckaccountHoldOperation];
 171             [strongSelf.operationQueue addOperation: fulfillBlock];
 172
 173             return YES;
 174         }
 175         return NO;
 176     }]]);
 177
 178     OCMStub([self.mockContainer accountInfoWithCompletionHandler:
 179              [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (CKAccountInfo* accountInfo,
 180                                                                NSError * error)) {
 181         __strong __typeof(self) strongSelf = weakSelf;
 182         if(passedBlock && strongSelf) {
 183             NSBlockOperation* fulfillBlock = [NSBlockOperation named:@"account-info-completion" withBlock: ^{
 184                 __strong __typeof(self) blockStrongSelf = weakSelf;
 185                 CKAccountInfo* account = [[CKAccountInfo alloc] init];
 186                 account.accountStatus = blockStrongSelf.accountStatus;
 187                 account.hasValidCredentials = blockStrongSelf.iCloudHasValidCredentials;
 188                 account.accountPartition = CKAccountPartitionTypeProduction;
 189                 passedBlock((CKAccountInfo*)account, nil);
 190             }];
 191             [fulfillBlock addDependency: strongSelf.ckaccountHoldOperation];
 192             [strongSelf.operationQueue addOperation: fulfillBlock];
 193
 194             return YES;
 195         }
 196         return NO;
 197     }]]);
 198
 199     self.mockAccountStateTracker = OCMClassMock([CKKSAccountStateTracker class]);
 200     OCMStub([self.mockAccountStateTracker getCircleStatus]).andCall(self, @selector(circleStatus));
 201
 202     // Fake out SOS peers
 203     // One trusted non-self peer, but it doesn't have any Octagon keys. Your test can change this if it wants.
 204     // However, note that [self putFakeDeviceStatusInCloudKit:] will likely not do what you want after you change this
 205     CKKSSOSSelfPeer* currentSelfPeer = [[CKKSSOSSelfPeer alloc] initWithSOSPeerID:@"local-peer"
 206                                                                     encryptionKey:[[SFECKeyPair alloc] initRandomKeyPairWithSpecifier:[[SFECKeySpecifier alloc] initWithCurve:SFEllipticCurveNistp384]]
 207                                                                        signingKey:[[SFECKeyPair alloc] initRandomKeyPairWithSpecifier:[[SFECKeySpecifier alloc] initWithCurve:SFEllipticCurveNistp384]]
 208                                                                          viewList:self.managedViewList];
 209
 210     self.mockSOSAdapter = [[CKKSMockSOSPresentAdapter alloc] initWithSelfPeer:currentSelfPeer
 211                                                                  trustedPeers:[NSSet set]
 212                                                                     essential:YES];
 213
 214     // If we're in circle, come up with a fake circle id. Otherwise, return an error.
 215     OCMStub([self.mockAccountStateTracker fetchCirclePeerID:
 216              [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (NSString* peerID,
 217                                                                NSError * error)) {
 218         __strong __typeof(self) strongSelf = weakSelf;
 219         if(passedBlock && strongSelf) {
 220             if(strongSelf.mockSOSAdapter.circleStatus == kSOSCCInCircle) {
 221                 passedBlock(strongSelf.mockSOSAdapter.selfPeer.peerID, nil);
 222             } else {
 223                 passedBlock(nil, [NSError errorWithDomain:@"securityd" code:errSecInternalError userInfo:@{NSLocalizedDescriptionKey:@"no account, no circle id"}]);
 224             }
 225
 226             return YES;
 227         }
 228         return NO;
 229     }]]);
 230
 231     self.aksLockState = false; // Lie and say AKS is always unlocked
 232     self.mockLockStateTracker = OCMClassMock([CKKSLockStateTracker class]);
 233     OCMStub([self.mockLockStateTracker queryAKSLocked]).andCall(self, @selector(aksLockState));
 234
 235     self.mockTTR = OCMClassMock([SecTapToRadar class]);
 236     OCMStub([self.mockTTR isRateLimited:[OCMArg any]]).andCall(self, @selector(isRateLimited:));
 237     OCMStub([self.mockTTR askUserIfTTR:[OCMArg any]]).andCall(self, @selector(askUserIfTTR:));
 238     OCMStub([self.mockTTR triggerTapToRadar:[OCMArg any]]).andCall(self, @selector(triggerTapToRadar:));
 239     self.isTTRRatelimited = true;
 240
 241     self.mockFakeCKModifyRecordZonesOperation = OCMClassMock([FakeCKModifyRecordZonesOperation class]);
 242     OCMStub([self.mockFakeCKModifyRecordZonesOperation ckdb]).andReturn(self.zones);
 243     OCMStub([self.mockFakeCKModifyRecordZonesOperation shouldFailModifyRecordZonesOperation]).andCall(self, @selector(shouldFailModifyRecordZonesOperation));
 244
 245     OCMStub([self.mockFakeCKModifyRecordZonesOperation ensureZoneDeletionAllowed:[OCMArg any]]).andCall(self, @selector(ensureZoneDeletionAllowed:));
 246
 247     self.mockFakeCKModifySubscriptionsOperation = OCMClassMock([FakeCKModifySubscriptionsOperation class]);
 248     OCMStub([self.mockFakeCKModifySubscriptionsOperation ckdb]).andReturn(self.zones);
 249
 250     self.mockFakeCKFetchRecordZoneChangesOperation = OCMClassMock([FakeCKFetchRecordZoneChangesOperation class]);
 251     OCMStub([self.mockFakeCKFetchRecordZoneChangesOperation ckdb]).andReturn(self.zones);
 252     OCMStub([self.mockFakeCKFetchRecordZoneChangesOperation isNetworkReachable]).andCall(self, @selector(isNetworkReachable));
 253
 254     self.mockFakeCKFetchRecordsOperation = OCMClassMock([FakeCKFetchRecordsOperation class]);
 255     OCMStub([self.mockFakeCKFetchRecordsOperation ckdb]).andReturn(self.zones);
 256
 257     self.mockFakeCKQueryOperation = OCMClassMock([FakeCKQueryOperation class]);
 258     OCMStub([self.mockFakeCKQueryOperation ckdb]).andReturn(self.zones);
 259
 260
 261     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 262         __strong __typeof(self) strongSelf = weakSelf;
 263         BOOL matches = NO;
 264         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 265             if(strongSelf.silentFetchesAllowed) {
 266                 matches = YES;
 267
 268                 FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 269                 [frzco addNullableDependency:strongSelf.ckFetchHoldOperation];
 270                 [strongSelf.operationQueue addOperation: frzco];
 271             }
 272         }
 273         return matches;
 274     }]]);
 275
 276     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 277         __strong __typeof(self) strongSelf = weakSelf;
 278         BOOL matches = NO;
 279         if ([obj isKindOfClass: [FakeCKFetchRecordsOperation class]]) {
 280             if(strongSelf.silentFetchesAllowed) {
 281                 matches = YES;
 282
 283                 FakeCKFetchRecordsOperation *ffro = (FakeCKFetchRecordsOperation *)obj;
 284                 [ffro addNullableDependency:strongSelf.ckFetchHoldOperation];
 285                 [strongSelf.operationQueue addOperation: ffro];
 286             }
 287         }
 288         return matches;
 289     }]]);
 290
 291     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 292         __strong __typeof(self) strongSelf = weakSelf;
 293         BOOL matches = NO;
 294         if ([obj isKindOfClass: [FakeCKQueryOperation class]]) {
 295             if(strongSelf.silentFetchesAllowed) {
 296                 matches = YES;
 297
 298                 FakeCKQueryOperation *fqo = (FakeCKQueryOperation *)obj;
 299                 [fqo addNullableDependency:strongSelf.ckFetchHoldOperation];
 300                 [strongSelf.operationQueue addOperation: fqo];
 301             }
 302         }
 303         return matches;
 304     }]]);
 305
 306     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 307         __strong __typeof(self) strongSelf = weakSelf;
 308         BOOL matches = NO;
 309         if ([obj isKindOfClass: [FakeCKModifyRecordZonesOperation class]]) {
 310             FakeCKModifyRecordZonesOperation *frzco = (FakeCKModifyRecordZonesOperation *)obj;
 311             [frzco addNullableDependency:strongSelf.ckModifyRecordZonesHoldOperation];
 312             [strongSelf.operationQueue addOperation: frzco];
 313             matches = YES;
 314         }
 315         return matches;
 316     }]]);
 317
 318     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 319         __strong __typeof(self) strongSelf = weakSelf;
 320         BOOL matches = NO;
 321         if ([obj isKindOfClass: [FakeCKModifySubscriptionsOperation class]]) {
 322             FakeCKModifySubscriptionsOperation *frzco = (FakeCKModifySubscriptionsOperation *)obj;
 323             [frzco addNullableDependency:strongSelf.ckModifySubscriptionsHoldOperation];
 324             [strongSelf.operationQueue addOperation: frzco];
 325             matches = YES;
 326         }
 327         return matches;
 328     }]]);
 329
 330     self.testZoneID = [[CKRecordZoneID alloc] initWithZoneName:@"testzone" ownerName:CKCurrentUserDefaultName];
 331
 332     // We don't want to use class mocks here, because they don't play well with partial mocks
 333     CKKSCloudKitClassDependencies* cloudKitClassDependencies = [[CKKSCloudKitClassDependencies alloc] initWithFetchRecordZoneChangesOperationClass:[FakeCKFetchRecordZoneChangesOperation class]
 334                                                                                                                         fetchRecordsOperationClass:[FakeCKFetchRecordsOperation class]
 335                                                                                                                                queryOperationClass:[FakeCKQueryOperation class]
 336                                                                                                                  modifySubscriptionsOperationClass:[FakeCKModifySubscriptionsOperation class]
 337                                                                                                                    modifyRecordZonesOperationClass:[FakeCKModifyRecordZonesOperation class]
 338                                                                                                                                 apsConnectionClass:[FakeAPSConnection class]
 339                                                                                                                          nsnotificationCenterClass:[FakeNSNotificationCenter class]
 340                                                                 nsdistributednotificationCenterClass:[FakeNSDistributedNotificationCenter class]
 341                                                                                                                                      notifierClass:[FakeCKKSNotifier class]];
 342
 343     self.mockCKKSViewManager = OCMPartialMock(
 344         [[CKKSViewManager alloc] initWithContainerName:SecCKKSContainerName
 345                                                 usePCS:SecCKKSContainerUsePCS
 346                                             sosAdapter:self.mockSOSAdapter
 347                              cloudKitClassDependencies:cloudKitClassDependencies]);
 348
 349     OCMStub([self.mockCKKSViewManager defaultViewList]).andCall(self, @selector(managedViewList));
 350     OCMStub([self.mockCKKSViewManager syncBackupAndNotifyAboutSync]);
 351     OCMStub([self.mockCKKSViewManager waitForTrustReady]).andReturn(YES);
 352
 353     self.injectedManager = self.mockCKKSViewManager;
 354
 355     [CKKSViewManager resetManager:false setTo:self.injectedManager];
 356
 357     // Lie and say network is available
 358     [self.reachabilityTracker setNetworkReachability:true];
 359
 360     // Make a new fake keychain
 361     NSString* tmp_dir = [NSString stringWithFormat: @"/tmp/%@.%X", testName, arc4random()];
 362     [[NSFileManager defaultManager] createDirectoryAtPath:[NSString stringWithFormat: @"%@/Library/Keychains", tmp_dir] withIntermediateDirectories:YES attributes:nil error:NULL];
 363
 364     SetCustomHomeURLString((__bridge CFStringRef) tmp_dir);
 365     SecKeychainDbReset(NULL);
 366
 367     // Actually load the database.
 368     kc_with_dbt(true, NULL, ^bool (SecDbConnectionRef dbt) { return false; });
 369 }
 370
 371 - (SOSAccountStatus*)circleStatus {
 372     NSError* error = nil;
 373     SOSCCStatus status = [self.mockSOSAdapter circleStatus:&error];
 374     return [[SOSAccountStatus alloc] init:status error:error];
 375 }
 376
 377 - (bool)aksLockState
 378 {
 379     return _aksLockState;
 380 }
 381
 382 - (void)setAksLockState:(bool)aksLockState
 383 {
 384
 385     if(aksLockState) {
 386         [SecMockAKS lockClassA];
 387     } else {
 388         [SecMockAKS unlockAllClasses];
 389     }
 390     _aksLockState = aksLockState;
 391 }
 392
 393 - (bool)isNetworkReachable {
 394     return self.reachabilityTracker.currentReachability;
 395 }
 396
 397 - (void)ckcontainerSubmitEventMetric:(CKEventMetric*)metric {
 398     @try {
 399         [self.mockContainerExpectations submitEventMetric:metric];
 400     } @catch (NSException *exception) {
 401         XCTFail("Received an container exception when trying to add a metric: %@", exception);
 402     }
 403 }
 404
 405 - (void)ckdatabaseAddOperation:(NSOperation*)op {
 406     @try {
 407         [self.mockDatabase addOperation:op];
 408     } @catch (NSException *exception) {
 409         XCTFail("Received an database exception: %@", exception);
 410     }
 411 }
 412
 413 - (NSError* _Nullable)shouldFailModifyRecordZonesOperation {
 414     NSError* error = self.nextModifyRecordZonesError;
 415     if(error) {
 416         self.nextModifyRecordZonesError = nil;
 417         return error;
 418     }
 419     return nil;
 420 }
 421
 422 - (void)ensureZoneDeletionAllowed:(FakeCKZone*)zone {
 423     XCTAssertTrue(self.silentZoneDeletesAllowed, "Should be allowing zone deletes");
 424 }
 425
 426 -(CKKSAccountStateTracker*)accountStateTracker {
 427     return self.injectedManager.accountTracker;
 428 }
 429
 430 -(CKKSLockStateTracker*)lockStateTracker {
 431     return self.injectedManager.lockStateTracker;
 432 }
 433
 434 -(CKKSReachabilityTracker*)reachabilityTracker {
 435     return self.injectedManager.reachabilityTracker;
 436 }
 437
 438 -(NSSet*)managedViewList {
 439     return (NSSet*) CFBridgingRelease(SOSViewCopyViewSet(kViewSetCKKS));
 440 }
 441
 442 -(void)expectCKFetch {
 443     [self expectCKFetchAndRunBeforeFinished: nil];
 444 }
 445
 446 -(void)expectCKFetchAndRunBeforeFinished: (void (^)(void))blockAfterFetch {
 447     [self expectCKFetchWithFilter:^BOOL(FakeCKFetchRecordZoneChangesOperation * op) {
 448         return YES;
 449     }
 450                 runBeforeFinished:blockAfterFetch];
 451 }
 452
 453 - (void)expectCKFetchWithFilter:(BOOL (^)(FakeCKFetchRecordZoneChangesOperation*))operationMatch
 454               runBeforeFinished:(void (^)(void))blockAfterFetch
 455 {
 456     // Create an object for the block to retain and modify
 457     BoolHolder* runAlready = [[BoolHolder alloc] init];
 458
 459     __weak __typeof(self) weakSelf = self;
 460     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 461         __strong __typeof(self) strongSelf = weakSelf;
 462         if(runAlready.state) {
 463             return NO;
 464         }
 465
 466         secnotice("fakecloudkit", "Received an operation (%@), checking if it's a fetch changes", obj);
 467         BOOL matches = NO;
 468         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 469             FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 470             matches = operationMatch(frzco);
 471             runAlready.state = true;
 472
 473             secnotice("fakecloudkit", "Running fetch changes: %@", obj);
 474             frzco.blockAfterFetch = blockAfterFetch;
 475             [frzco addNullableDependency: strongSelf.ckFetchHoldOperation];
 476             [strongSelf.operationQueue addOperation: frzco];
 477         }
 478         return matches;
 479     }]];
 480 }
 481
 482 -(void)expectCKFetchByRecordID {
 483     // Create an object for the block to retain and modify
 484     BoolHolder* runAlready = [[BoolHolder alloc] init];
 485
 486     __weak __typeof(self) weakSelf = self;
 487     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 488         __strong __typeof(self) strongSelf = weakSelf;
 489         if(runAlready.state) {
 490             return NO;
 491         }
 492         BOOL matches = NO;
 493         if ([obj isKindOfClass: [FakeCKFetchRecordsOperation class]]) {
 494             matches = YES;
 495             runAlready.state = true;
 496
 497             FakeCKFetchRecordsOperation *ffro = (FakeCKFetchRecordsOperation *)obj;
 498             [ffro addNullableDependency: strongSelf.ckFetchHoldOperation];
 499             [strongSelf.operationQueue addOperation: ffro];
 500         }
 501         return matches;
 502     }]];
 503 }
 504
 505
 506 -(void)expectCKFetchByQuery {
 507     // Create an object for the block to retain and modify
 508     BoolHolder* runAlready = [[BoolHolder alloc] init];
 509
 510     __weak __typeof(self) weakSelf = self;
 511     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 512         __strong __typeof(self) strongSelf = weakSelf;
 513         if(runAlready.state) {
 514             return NO;
 515         }
 516         BOOL matches = NO;
 517         if ([obj isKindOfClass: [FakeCKQueryOperation class]]) {
 518             matches = YES;
 519             runAlready.state = true;
 520
 521             FakeCKQueryOperation *fqo = (FakeCKQueryOperation *)obj;
 522             [fqo addNullableDependency: strongSelf.ckFetchHoldOperation];
 523             [strongSelf.operationQueue addOperation: fqo];
 524         }
 525         return matches;
 526     }]];
 527 }
 528
 529 - (void)startCKKSSubsystem {
 530     if(self.fakeHSA2AccountStatus != CKKSAccountStatusUnknown) {
 531         [self.accountStateTracker setHSA2iCloudAccountStatus:self.fakeHSA2AccountStatus];
 532     }
 533     [self startCKAccountStatusMock];
 534 }
 535
 536 - (void)startCKAccountStatusMock {
 537     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 538     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 539     // and return; it's just a tad confusing.
 540     if([self.ckaccountHoldOperation isPending]) {
 541         [self.operationQueue addOperation: self.ckaccountHoldOperation];
 542     }
 543
 544     [self.accountStateTracker performInitialDispatches];
 545 }
 546
 547 -(void)holdCloudKitModifications {
 548     XCTAssertFalse([self.ckModifyHoldOperation isPending], "Shouldn't already be a pending cloudkit modify hold operation");
 549     self.ckModifyHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 550         secnotice("ckks", "Released CloudKit modification hold.");
 551     }];
 552 }
 553 -(void)releaseCloudKitModificationHold {
 554     if([self.ckModifyHoldOperation isPending]) {
 555         [self.operationQueue addOperation: self.ckModifyHoldOperation];
 556     }
 557 }
 558
 559 -(void)holdCloudKitFetches {
 560     XCTAssertFalse([self.ckFetchHoldOperation isPending], "Shouldn't already be a pending cloudkit fetch hold operation");
 561     self.ckFetchHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 562         secnotice("ckks", "Released CloudKit fetch hold.");
 563     }];
 564 }
 565 -(void)releaseCloudKitFetchHold {
 566     if([self.ckFetchHoldOperation isPending]) {
 567         [self.operationQueue addOperation: self.ckFetchHoldOperation];
 568     }
 569 }
 570
 571 -(void)holdCloudKitModifyRecordZones {
 572     XCTAssertFalse([self.ckModifyRecordZonesHoldOperation isPending], "Shouldn't already be a pending cloudkit zone create hold operation");
 573     self.ckModifyRecordZonesHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 574         secnotice("ckks", "Released CloudKit zone create hold.");
 575     }];
 576 }
 577 -(void)releaseCloudKitModifyRecordZonesHold {
 578     if([self.ckModifyRecordZonesHoldOperation isPending]) {
 579         [self.operationQueue addOperation: self.ckModifyRecordZonesHoldOperation];
 580     }
 581 }
 582
 583 -(void)holdCloudKitModifySubscription {
 584     XCTAssertFalse([self.ckModifySubscriptionsHoldOperation isPending], "Shouldn't already be a pending cloudkit subscription hold operation");
 585     self.ckModifySubscriptionsHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 586         secnotice("ckks", "Released CloudKit zone create hold.");
 587     }];
 588 }
 589 -(void)releaseCloudKitModifySubscriptionHold {
 590     if([self.ckModifySubscriptionsHoldOperation isPending]) {
 591         [self.operationQueue addOperation: self.ckModifySubscriptionsHoldOperation];
 592     }
 593 }
 594
 595 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID {
 596     [self expectCKModifyItemRecords:expectedNumberOfRecords
 597            currentKeyPointerRecords:expectedCurrentKeyRecords
 598                              zoneID:zoneID
 599                           checkItem:nil];
 600 }
 601
 602 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID checkItem: (BOOL (^)(CKRecord*)) checkItem {
 603     [self expectCKModifyItemRecords:expectedNumberOfRecords
 604                      deletedRecords:0
 605            currentKeyPointerRecords:expectedCurrentKeyRecords
 606                              zoneID:zoneID
 607                           checkItem:checkItem];
 608 }
 609
 610 - (void)expectCKModifyItemRecords:(NSUInteger)expectedNumberOfModifiedRecords
 611                    deletedRecords:(NSUInteger)expectedNumberOfDeletedRecords
 612          currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 613                            zoneID:(CKRecordZoneID*)zoneID
 614                         checkItem:(BOOL (^)(CKRecord*))checkItem {
 615     // We're updating the device state type on every update, so add it in here
 616     NSMutableDictionary* expectedRecords = [@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfModifiedRecords],
 617                                               SecCKRecordCurrentKeyType: [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords],
 618                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 619                                               } mutableCopy];
 620
 621     if(SecCKKSSyncManifests()) {
 622         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 1];
 623         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 624     }
 625
 626     NSDictionary* deletedRecords = nil;
 627     if(expectedNumberOfDeletedRecords != 0) {
 628         deletedRecords = @{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfDeletedRecords]};
 629     }
 630
 631     [self expectCKModifyRecords:expectedRecords
 632         deletedRecordTypeCounts:deletedRecords
 633                          zoneID:zoneID
 634             checkModifiedRecord: ^BOOL (CKRecord* record){
 635                 if([record.recordType isEqualToString: SecCKRecordItemType] && checkItem) {
 636                     return checkItem(record);
 637                 } else {
 638                     return YES;
 639                 }
 640             }
 641            runAfterModification:nil];
 642 }
 643
 644
 645
 646 - (void)expectCKModifyKeyRecords:(NSUInteger)expectedNumberOfRecords
 647         currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 648                  tlkShareRecords:(NSUInteger)expectedTLKShareRecords
 649                           zoneID:(CKRecordZoneID*)zoneID
 650 {
 651     return [self expectCKModifyKeyRecords:expectedNumberOfRecords
 652                  currentKeyPointerRecords:expectedCurrentKeyRecords
 653                           tlkShareRecords:expectedTLKShareRecords
 654                                    zoneID:zoneID
 655                       checkModifiedRecord:nil];
 656 }
 657
 658 - (void)expectCKModifyKeyRecords:(NSUInteger)expectedNumberOfRecords
 659         currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 660                  tlkShareRecords:(NSUInteger)expectedTLKShareRecords
 661                           zoneID:(CKRecordZoneID*)zoneID
 662              checkModifiedRecord:(BOOL (^_Nullable)(CKRecord*))checkModifiedRecord
 663 {
 664     NSNumber* nkeys = [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords];
 665     NSNumber* ncurrentkeys = [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords];
 666     NSNumber* ntlkshares = [NSNumber numberWithUnsignedInteger: expectedTLKShareRecords];
 667
 668     [self expectCKModifyRecords:@{SecCKRecordIntermediateKeyType: nkeys,
 669                                   SecCKRecordCurrentKeyType: ncurrentkeys,
 670                                   SecCKRecordTLKShareType: ntlkshares,
 671                                   }
 672         deletedRecordTypeCounts:nil
 673                          zoneID:zoneID
 674             checkModifiedRecord:checkModifiedRecord
 675            runAfterModification:nil];
 676 }
 677
 678 - (void)expectCKModifyRecords:(NSDictionary<NSString*, NSNumber*>*) expectedRecordTypeCounts
 679       deletedRecordTypeCounts:(NSDictionary<NSString*, NSNumber*>*) expectedDeletedRecordTypeCounts
 680                        zoneID:(CKRecordZoneID*) zoneID
 681           checkModifiedRecord:(BOOL (^)(CKRecord*)) checkModifiedRecord
 682          runAfterModification:(void (^) (void))afterModification
 683 {
 684     __weak __typeof(self) weakSelf = self;
 685
 686     // Create an object for the block to retain and modify
 687     BoolHolder* runAlready = [[BoolHolder alloc] init];
 688
 689     secnotice("fakecloudkit", "expecting an operation matching modifications: %@ deletions: %@",
 690               expectedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 691
 692     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 693         secnotice("fakecloudkit", "Received an operation (%@), checking if it's a modification", obj);
 694         __block bool matches = false;
 695         if(runAlready.state) {
 696             secnotice("fakecloudkit", "Run already, skipping");
 697             return NO;
 698         }
 699
 700         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 701             __strong __typeof(weakSelf) strongSelf = weakSelf;
 702             XCTAssertNotNil(strongSelf, "self exists");
 703
 704             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 705             matches = true;
 706
 707             NSMutableDictionary<NSString*, NSNumber*>* modifiedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 708             NSMutableDictionary<NSString*, NSNumber*>* deletedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 709
 710             // First: check if it matches. If it does, _then_ execute the operation.
 711             // Supports single-zone atomic writes only
 712
 713             if(!op.atomic) {
 714                 // We only care about atomic operations
 715                 secnotice("fakecloudkit", "Not an atomic operation; quitting: %@", op);
 716                 return NO;
 717             }
 718
 719             FakeCKZone* zone = strongSelf.zones[zoneID];
 720             XCTAssertNotNil(zone, "Have a zone for these records");
 721
 722             __block BOOL result = YES;
 723             dispatch_sync(zone.queue, ^{
 724
 725                 for(CKRecord* record in op.recordsToSave) {
 726                     if(![record.recordID.zoneID isEqual: zoneID]) {
 727                         secnotice("fakecloudkit", "Modified record zone ID mismatch: %@ %@", zoneID, record.recordID.zoneID);
 728                         result = NO;
 729                         return;
 730                     }
 731
 732                     NSError* recordError = [zone errorFromSavingRecord: record];
 733                     if(recordError) {
 734                         secnotice("fakecloudkit", "Record zone rejected record write: %@ %@", recordError, record);
 735                         XCTFail(@"Record zone rejected record write: %@ %@", recordError, record);
 736                         result = NO;
 737                         return;
 738                     }
 739
 740                     NSNumber* currentCountNumber = modifiedRecordTypeCounts[record.recordType];
 741                     NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 742                     modifiedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 743                 }
 744
 745                 for(CKRecordID* recordID in op.recordIDsToDelete) {
 746                     if(![recordID.zoneID isEqual: zoneID]) {
 747                         matches = false;
 748                         secnotice("fakecloudkit", "Deleted record zone ID mismatch: %@ %@", zoneID, recordID.zoneID);
 749                     }
 750
 751                     // Find the object in CloudKit, and record its type
 752                     CKRecord* record = strongSelf.zones[zoneID].currentDatabase[recordID];
 753                     if(record) {
 754                         NSNumber* currentCountNumber = deletedRecordTypeCounts[record.recordType];
 755                         NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 756                         deletedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 757                     }
 758                 }
 759
 760                 NSMutableDictionary* filteredExpectedRecordTypeCounts = [expectedRecordTypeCounts mutableCopy];
 761                 for(NSString* key in filteredExpectedRecordTypeCounts.allKeys) {
 762                     if([filteredExpectedRecordTypeCounts[key] isEqual: [NSNumber numberWithInt:0]]) {
 763                         filteredExpectedRecordTypeCounts[key] = nil;
 764                     }
 765                 }
 766                 filteredExpectedRecordTypeCounts[SecCKRecordManifestType] = modifiedRecordTypeCounts[SecCKRecordManifestType];
 767                 filteredExpectedRecordTypeCounts[SecCKRecordManifestLeafType] = modifiedRecordTypeCounts[SecCKRecordManifestLeafType];
 768
 769                 // Inspect that we have exactly the same records as we expect
 770                 if(expectedRecordTypeCounts) {
 771                     matches &= !![modifiedRecordTypeCounts isEqual: filteredExpectedRecordTypeCounts];
 772                     if(!matches) {
 773                         secnotice("fakecloudkit", "Record number mismatch: %@ %@", modifiedRecordTypeCounts, filteredExpectedRecordTypeCounts);
 774                         result = NO;
 775                         return;
 776                     }
 777                 } else {
 778                     matches &= op.recordsToSave.count == 0u;
 779                     if(!matches) {
 780                         secnotice("fakecloudkit", "Record number mismatch: %@ 0", modifiedRecordTypeCounts);
 781                         result = NO;
 782                         return;
 783                     }
 784                 }
 785                 if(expectedDeletedRecordTypeCounts) {
 786                     matches &= !![deletedRecordTypeCounts  isEqual: expectedDeletedRecordTypeCounts];
 787                     if(!matches) {
 788                         secnotice("fakecloudkit", "Deleted record number mismatch: %@ %@", deletedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 789                         result = NO;
 790                         return;
 791                     }
 792                 } else {
 793                     matches &= op.recordIDsToDelete.count == 0u;
 794                     if(!matches) {
 795                         secnotice("fakecloudkit", "Deleted record number mismatch: %@ 0", deletedRecordTypeCounts);
 796                         result = NO;
 797                         return;
 798                     }
 799                 }
 800
 801                 // We have the right number of things, and their etags match. Ensure that they have the right etags
 802                 if(matches && checkModifiedRecord) {
 803                     // Clearly we have the right number of things. Call checkRecord on them...
 804                     for(CKRecord* record in op.recordsToSave) {
 805                         matches &= !!(checkModifiedRecord(record));
 806                         if(!matches) {
 807                             secnotice("fakecloudkit", "Check record reports NO: %@ 0", record);
 808                             result = NO;
 809                             return;
 810                         }
 811                     }
 812                 }
 813
 814                 if(matches) {
 815                     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 816                     // if you'd like to read the data from this write.
 817                     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-write" withBlock: ^{
 818                         @synchronized(zone.currentDatabase) {
 819                             NSMutableArray* savedRecords = [[NSMutableArray alloc] init];
 820                             for(CKRecord* record in op.recordsToSave) {
 821                                 CKRecord* reflectedRecord = [record copy];
 822                                 reflectedRecord.modificationDate = [NSDate date];
 823
 824                                 [zone addToZone: reflectedRecord];
 825
 826                                 [savedRecords addObject:reflectedRecord];
 827                                 op.perRecordCompletionBlock(reflectedRecord, nil);
 828                             }
 829                             for(CKRecordID* recordID in op.recordIDsToDelete) {
 830                                 // I don't believe CloudKit fails an operation if you delete a record that's not there, so:
 831                                 [zone deleteCKRecordIDFromZone: recordID];
 832                             }
 833
 834                             if(afterModification) {
 835                                 afterModification();
 836                             }
 837
 838                             op.modifyRecordsCompletionBlock(savedRecords, op.recordIDsToDelete, nil);
 839                             op.isFinished = YES;
 840                         }
 841                     }];
 842                     [ckop addNullableDependency:strongSelf.ckModifyHoldOperation];
 843                     [strongSelf.operationQueue addOperation: ckop];
 844                 }
 845             });
 846             if(result != YES) {
 847                 return result;
 848             }
 849         }
 850         if(matches) {
 851             runAlready.state = true;
 852         }
 853         return matches ? YES : NO;
 854     }]];
 855 }
 856
 857 - (void)failNextZoneCreation:(CKRecordZoneID*)zoneID {
 858     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 859     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 860     self.zones[zoneID].creationError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain
 861                                                                         code:CKErrorNetworkUnavailable
 862                                                                     userInfo:@{
 863                                                                                CKErrorRetryAfterKey: @(0.5),
 864                                                                                }];
 865 }
 866
 867 // Report success, but don't actually create the zone.
 868 // This way, you can find ZoneNotFound errors later on
 869 - (void)failNextZoneCreationSilently:(CKRecordZoneID*)zoneID {
 870     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 871     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 872     self.zones[zoneID].failCreationSilently = true;
 873 }
 874
 875 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID {
 876     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 877     self.zones[zoneID].subscriptionError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 878 }
 879
 880 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID withError:(NSError*)error {
 881     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 882     self.zones[zoneID].subscriptionError = error;
 883 }
 884
 885 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 886     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:nil];
 887 }
 888
 889 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)(void))blockAfterReject {
 890     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:blockAfterReject withError:nil];
 891 }
 892
 893 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)(void))blockAfterReject withError:(NSError*)error {
 894     __weak __typeof(self) weakSelf = self;
 895
 896     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 897         __strong __typeof(weakSelf) strongSelf = weakSelf;
 898         XCTAssertNotNil(strongSelf, "self exists");
 899
 900         __block bool rejected = false;
 901         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 902             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 903
 904             if(!op.atomic) {
 905                 // We only care about atomic operations
 906                 return NO;
 907             }
 908
 909             // We want to only match zone updates pertaining to this zone
 910             for(CKRecord* record in op.recordsToSave) {
 911                 if(![record.recordID.zoneID isEqual: zoneID]) {
 912                     return NO;
 913                 }
 914             }
 915
 916             FakeCKZone* zone = strongSelf.zones[zoneID];
 917             XCTAssertNotNil(zone, "Have a zone for these records");
 918
 919             rejected = true;
 920
 921             if(error) {
 922                 [strongSelf rejectWrite: op withError:error];
 923             } else {
 924                 NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 925                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 926             }
 927
 928             if(blockAfterReject) {
 929                 blockAfterReject();
 930             }
 931         }
 932         return rejected ? YES : NO;
 933     }]];
 934 }
 935
 936 - (void)expectCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 937     __weak __typeof(self) weakSelf = self;
 938
 939     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 940         __strong __typeof(weakSelf) strongSelf = weakSelf;
 941         XCTAssertNotNil(strongSelf, "self exists");
 942
 943         __block bool rejected = false;
 944         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 945             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 946
 947             secnotice("fakecloudkit", "checking for expectCKAtomicModifyItemRecordsUpdateFailure");
 948
 949             if(!op.atomic) {
 950                 // We only care about atomic operations
 951                 secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: update not atomic");
 952                 return NO;
 953             }
 954
 955             // We want to only match zone updates pertaining to this zone
 956             for(CKRecord* record in op.recordsToSave) {
 957                 if(![record.recordID.zoneID isEqual: zoneID]) {
 958                     secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: %@ is not %@", record.recordID.zoneID, zoneID);
 959                     return NO;
 960                 }
 961             }
 962
 963             FakeCKZone* zone = strongSelf.zones[zoneID];
 964             XCTAssertNotNil(zone, "Have a zone for these records");
 965
 966             NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 967
 968             @synchronized(zone.currentDatabase) {
 969                 for(CKRecord* record in op.recordsToSave) {
 970                     // Check if we should allow this transaction
 971                     NSError* recordSaveError = [zone errorFromSavingRecord: record];
 972                     if(recordSaveError) {
 973                         failedRecords[record.recordID] = recordSaveError;
 974                         rejected = true;
 975                     }
 976                 }
 977             }
 978
 979             if(rejected) {
 980                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 981             } else {
 982                 secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: doesn't seem like an error to us");
 983             }
 984         }
 985         return rejected ? YES : NO;
 986     }]];
 987 }
 988
 989 -(void)rejectWrite:(CKModifyRecordsOperation*)op withError:(NSError*)error {
 990     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 991     // if you'd like to read the data from this write.
 992     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write-error" withBlock: ^{
 993         op.modifyRecordsCompletionBlock(nil, nil, error);
 994         op.isFinished = YES;
 995     }];
 996     [ckop addNullableDependency: self.ckModifyHoldOperation];
 997     [self.operationQueue addOperation: ckop];
 998 }
 999
1000 -(void)rejectWrite:(CKModifyRecordsOperation*)op failedRecords:(NSMutableDictionary<CKRecordID*, NSError*>*)failedRecords {
1001     // Add the batch request failed errors
1002     for(CKRecord* record in op.recordsToSave) {
1003         NSError* exists = failedRecords[record.recordID];
1004         if(!exists) {
1005             // TODO: might have important userInfo, but we're not mocking that yet
1006             failedRecords[record.recordID] = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorBatchRequestFailed userInfo: @{}];
1007         }
1008     }
1009
1010     NSError* error = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorPartialFailure userInfo: @{CKPartialErrorsByItemIDKey: failedRecords}];
1011
1012     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
1013     // if you'd like to read the data from this write.
1014     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write" withBlock: ^{
1015         op.modifyRecordsCompletionBlock(nil, nil, error);
1016         op.isFinished = YES;
1017     }];
1018     [ckop addNullableDependency: self.ckModifyHoldOperation];
1019     [self.operationQueue addOperation: ckop];
1020 }
1021
1022 - (void)expectCKDeleteItemRecords:(NSUInteger)expectedNumberOfRecords
1023                            zoneID:(CKRecordZoneID*) zoneID {
1024
1025     // We're updating the device state type on every update, so add it in here
1026     NSMutableDictionary* expectedRecords = [@{
1027                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInteger:expectedNumberOfRecords],
1028                                               } mutableCopy];
1029     if(SecCKKSSyncManifests()) {
1030         // TODO: this really shouldn't be 2.
1031         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 2];
1032         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
1033     }
1034
1035     [self expectCKModifyRecords:expectedRecords
1036         deletedRecordTypeCounts:@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords]}
1037                          zoneID:zoneID
1038             checkModifiedRecord:nil
1039            runAfterModification:nil];
1040 }
1041
1042 -(void)waitForCKModifications {
1043     // CloudKit modifications are put on the local queue.
1044     // This is heavyweight but should suffice.
1045     [self.operationQueue waitUntilAllOperationsAreFinished];
1046 }
1047
1048 - (void)tearDown {
1049     NSString* testName = [self.name componentsSeparatedByString:@" "][1];
1050     testName = [testName stringByReplacingOccurrencesOfString:@"]" withString:@""];
1051     secnotice("ckkstest", "Ending test %@", testName);
1052
1053     if(SecCKKSIsEnabled()) {
1054         self.accountStatus = CKAccountStatusCouldNotDetermine;
1055
1056         // If the test never initialized the account state, don't call status later
1057         bool callStatus = [self.ckaccountHoldOperation isFinished];
1058         [self.ckaccountHoldOperation cancel];
1059         self.ckaccountHoldOperation = nil;
1060
1061         // Ensure we don't have any blocking operations left
1062         [self.operationQueue cancelAllOperations];
1063         [self waitForCKModifications];
1064
1065         XCTAssertEqual(0, [self.injectedManager.completedSecCKKSInitialize wait:20*NSEC_PER_SEC],
1066             "Timeout did not occur waiting for SecCKKSInitialize");
1067
1068         // Ensure that we can fetch zone status for all zones
1069         if(callStatus) {
1070             XCTestExpectation *statusReturned = [self expectationWithDescription:@"status returned"];
1071             [self.injectedManager rpcStatus:nil reply:^(NSArray<NSDictionary *> *result, NSError *error) {
1072                 XCTAssertNil(error, "Should be no error fetching status");
1073                 [statusReturned fulfill];
1074             }];
1075             [self waitForExpectations: @[statusReturned] timeout:20];
1076
1077             // Make sure this happens before teardown.
1078             XCTAssertEqual(0, [self.accountStateTracker.finishedInitialDispatches wait:20*NSEC_PER_SEC], "Account state tracker initialized itself");
1079
1080             dispatch_group_t accountChangesDelivered = [self.accountStateTracker checkForAllDeliveries];
1081             XCTAssertEqual(0, dispatch_group_wait(accountChangesDelivered, dispatch_time(DISPATCH_TIME_NOW, 10*NSEC_PER_SEC)), "Account state tracker finished delivering everything");
1082         }
1083     }
1084
1085     [super tearDown];
1086
1087     [self.injectedManager cancelPendingOperations];
1088     [CKKSViewManager resetManager:true setTo:nil];
1089     self.injectedManager = nil;
1090     [self.mockCKKSViewManager stopMocking];
1091     self.mockCKKSViewManager = nil;
1092
1093     [self.mockAccountStateTracker stopMocking];
1094     self.mockAccountStateTracker = nil;
1095
1096     [self.mockLockStateTracker stopMocking];
1097     self.mockLockStateTracker = nil;
1098
1099     [self.mockFakeCKModifyRecordZonesOperation stopMocking];
1100     self.mockFakeCKModifyRecordZonesOperation = nil;
1101
1102     [self.mockFakeCKModifySubscriptionsOperation stopMocking];
1103     self.mockFakeCKModifySubscriptionsOperation = nil;
1104
1105     [self.mockFakeCKFetchRecordZoneChangesOperation stopMocking];
1106     self.mockFakeCKFetchRecordZoneChangesOperation = nil;
1107
1108     [self.mockFakeCKFetchRecordsOperation stopMocking];
1109     self.mockFakeCKFetchRecordsOperation = nil;
1110
1111     [self.mockFakeCKQueryOperation stopMocking];
1112     self.mockFakeCKQueryOperation = nil;
1113
1114     [self.mockDatabase stopMocking];
1115     self.mockDatabase = nil;
1116
1117     [self.mockDatabaseExceptionCatcher stopMocking];
1118     self.mockDatabaseExceptionCatcher = nil;
1119
1120     [self.mockContainer stopMocking];
1121     self.mockContainer = nil;
1122
1123     [self.mockTTR stopMocking];
1124     self.mockTTR = nil;
1125     self.ttrExpectation = nil;
1126     self.isTTRRatelimited = true;
1127
1128     self.zones = nil;
1129
1130     _mockSOSAdapter = nil;
1131     _mockOctagonAdapter = nil;
1132
1133     SecCKKSTestResetFlags();
1134 }
1135
1136 - (CKKSKey*) fakeTLK: (CKRecordZoneID*)zoneID {
1137     CKKSKey* key = [[CKKSKey alloc] initSelfWrappedWithAESKey:[[CKKSAESSIVKey alloc] initWithBase64: @"uImdbZ7Zg+6WJXScTnRBfNmoU1UiMkSYxWc+d1Vuq3IFn2RmTRkTdWTe3HmeWo1pAomqy+upK8KHg2PGiRGhqg=="]
1138                                                          uuid:[[NSUUID UUID] UUIDString]
1139                                                      keyclass:SecCKKSKeyClassTLK
1140                                                         state: SecCKKSProcessedStateLocal
1141                                                        zoneID:zoneID
1142                                               encodedCKRecord: nil
1143                                                    currentkey: true];
1144     [key CKRecordWithZoneID: zoneID];
1145     return key;
1146 }
1147
1148 - (NSError*)ckInternalServerExtensionError:(NSInteger)code description:(NSString*)desc {
1149     return [FakeCKZone internalPluginError:@"CloudkitKeychainService" code:code description:desc];
1150 }
1151
1152 @end
1153
1154 #endif