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     self.testZoneID = [[CKRecordZoneID alloc] initWithZoneName:@"testzone" ownerName:CKCurrentUserDefaultName];
 307
 308     // We don't want to use class mocks here, because they don't play well with partial mocks
 309     CKKSCloudKitClassDependencies* cloudKitClassDependencies = [[CKKSCloudKitClassDependencies alloc] initWithFetchRecordZoneChangesOperationClass:[FakeCKFetchRecordZoneChangesOperation class]
 310                                                                                                                         fetchRecordsOperationClass:[FakeCKFetchRecordsOperation class]
 311                                                                                                                                queryOperationClass:[FakeCKQueryOperation class]
 312                                                                                                                  modifySubscriptionsOperationClass:[FakeCKModifySubscriptionsOperation class]
 313                                                                                                                    modifyRecordZonesOperationClass:[FakeCKModifyRecordZonesOperation class]
 314                                                                                                                                 apsConnectionClass:[FakeAPSConnection class]
 315                                                                                                                          nsnotificationCenterClass:[FakeNSNotificationCenter class]
 316                                                                 nsdistributednotificationCenterClass:[FakeNSDistributedNotificationCenter class]
 317                                                                                                                                      notifierClass:[FakeCKKSNotifier class]];
 318
 319     self.mockCKKSViewManager = OCMPartialMock(
 320         [[CKKSViewManager alloc] initWithContainerName:SecCKKSContainerName
 321                                                 usePCS:SecCKKSContainerUsePCS
 322                                             sosAdapter:self.mockSOSAdapter
 323                              cloudKitClassDependencies:cloudKitClassDependencies]);
 324
 325     OCMStub([self.mockCKKSViewManager defaultViewList]).andCall(self, @selector(managedViewList));
 326     OCMStub([self.mockCKKSViewManager syncBackupAndNotifyAboutSync]);
 327     OCMStub([self.mockCKKSViewManager waitForTrustReady]).andReturn(YES);
 328
 329     self.injectedManager = self.mockCKKSViewManager;
 330
 331     [CKKSViewManager resetManager:false setTo:self.injectedManager];
 332
 333     // Lie and say network is available
 334     [self.reachabilityTracker setNetworkReachability:true];
 335
 336     // Make a new fake keychain
 337     NSString* tmp_dir = [NSString stringWithFormat: @"/tmp/%@.%X", testName, arc4random()];
 338     [[NSFileManager defaultManager] createDirectoryAtPath:[NSString stringWithFormat: @"%@/Library/Keychains", tmp_dir] withIntermediateDirectories:YES attributes:nil error:NULL];
 339
 340     SetCustomHomeURLString((__bridge CFStringRef) tmp_dir);
 341     SecKeychainDbReset(NULL);
 342
 343     // Actually load the database.
 344     kc_with_dbt(true, NULL, ^bool (SecDbConnectionRef dbt) { return false; });
 345 }
 346
 347 - (SOSAccountStatus*)circleStatus {
 348     NSError* error = nil;
 349     SOSCCStatus status = [self.mockSOSAdapter circleStatus:&error];
 350     return [[SOSAccountStatus alloc] init:status error:error];
 351 }
 352
 353 - (bool)aksLockState
 354 {
 355     return _aksLockState;
 356 }
 357
 358 - (void)setAksLockState:(bool)aksLockState
 359 {
 360
 361     if(aksLockState) {
 362         [SecMockAKS lockClassA];
 363     } else {
 364         [SecMockAKS unlockAllClasses];
 365     }
 366     _aksLockState = aksLockState;
 367 }
 368
 369 - (bool)isNetworkReachable {
 370     return self.reachabilityTracker.currentReachability;
 371 }
 372
 373 - (void)ckcontainerSubmitEventMetric:(CKEventMetric*)metric {
 374     @try {
 375         [self.mockContainerExpectations submitEventMetric:metric];
 376     } @catch (NSException *exception) {
 377         XCTFail("Received an container exception when trying to add a metric: %@", exception);
 378     }
 379 }
 380
 381 - (void)ckdatabaseAddOperation:(NSOperation*)op {
 382     @try {
 383         [self.mockDatabase addOperation:op];
 384     } @catch (NSException *exception) {
 385         XCTFail("Received an database exception: %@", exception);
 386     }
 387 }
 388
 389 - (NSError* _Nullable)shouldFailModifyRecordZonesOperation {
 390     NSError* error = self.nextModifyRecordZonesError;
 391     if(error) {
 392         self.nextModifyRecordZonesError = nil;
 393         return error;
 394     }
 395     return nil;
 396 }
 397
 398 - (void)ensureZoneDeletionAllowed:(FakeCKZone*)zone {
 399     XCTAssertTrue(self.silentZoneDeletesAllowed, "Should be allowing zone deletes");
 400 }
 401
 402 -(CKKSAccountStateTracker*)accountStateTracker {
 403     return self.injectedManager.accountTracker;
 404 }
 405
 406 -(CKKSLockStateTracker*)lockStateTracker {
 407     return self.injectedManager.lockStateTracker;
 408 }
 409
 410 -(CKKSReachabilityTracker*)reachabilityTracker {
 411     return self.injectedManager.reachabilityTracker;
 412 }
 413
 414 -(NSSet*)managedViewList {
 415     return (NSSet*) CFBridgingRelease(SOSViewCopyViewSet(kViewSetCKKS));
 416 }
 417
 418 -(void)expectCKFetch {
 419     [self expectCKFetchAndRunBeforeFinished: nil];
 420 }
 421
 422 -(void)expectCKFetchAndRunBeforeFinished: (void (^)(void))blockAfterFetch {
 423     [self expectCKFetchWithFilter:^BOOL(FakeCKFetchRecordZoneChangesOperation * op) {
 424         return YES;
 425     }
 426                 runBeforeFinished:blockAfterFetch];
 427 }
 428
 429 - (void)expectCKFetchWithFilter:(BOOL (^)(FakeCKFetchRecordZoneChangesOperation*))operationMatch
 430               runBeforeFinished:(void (^)(void))blockAfterFetch
 431 {
 432     // Create an object for the block to retain and modify
 433     BoolHolder* runAlready = [[BoolHolder alloc] init];
 434
 435     __weak __typeof(self) weakSelf = self;
 436     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 437         __strong __typeof(self) strongSelf = weakSelf;
 438         if(runAlready.state) {
 439             return NO;
 440         }
 441
 442         secnotice("fakecloudkit", "Received an operation (%@), checking if it's a fetch changes", obj);
 443         BOOL matches = NO;
 444         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 445             FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 446             matches = operationMatch(frzco);
 447             runAlready.state = true;
 448
 449             secnotice("fakecloudkit", "Running fetch changes: %@", obj);
 450             frzco.blockAfterFetch = blockAfterFetch;
 451             [frzco addNullableDependency: strongSelf.ckFetchHoldOperation];
 452             [strongSelf.operationQueue addOperation: frzco];
 453         }
 454         return matches;
 455     }]];
 456 }
 457
 458 -(void)expectCKFetchByRecordID {
 459     // Create an object for the block to retain and modify
 460     BoolHolder* runAlready = [[BoolHolder alloc] init];
 461
 462     __weak __typeof(self) weakSelf = self;
 463     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 464         __strong __typeof(self) strongSelf = weakSelf;
 465         if(runAlready.state) {
 466             return NO;
 467         }
 468         BOOL matches = NO;
 469         if ([obj isKindOfClass: [FakeCKFetchRecordsOperation class]]) {
 470             matches = YES;
 471             runAlready.state = true;
 472
 473             FakeCKFetchRecordsOperation *ffro = (FakeCKFetchRecordsOperation *)obj;
 474             [ffro addNullableDependency: strongSelf.ckFetchHoldOperation];
 475             [strongSelf.operationQueue addOperation: ffro];
 476         }
 477         return matches;
 478     }]];
 479 }
 480
 481
 482 -(void)expectCKFetchByQuery {
 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: [FakeCKQueryOperation class]]) {
 494             matches = YES;
 495             runAlready.state = true;
 496
 497             FakeCKQueryOperation *fqo = (FakeCKQueryOperation *)obj;
 498             [fqo addNullableDependency: strongSelf.ckFetchHoldOperation];
 499             [strongSelf.operationQueue addOperation: fqo];
 500         }
 501         return matches;
 502     }]];
 503 }
 504
 505 - (void)startCKKSSubsystem {
 506     if(self.fakeHSA2AccountStatus != CKKSAccountStatusUnknown) {
 507         [self.accountStateTracker setHSA2iCloudAccountStatus:self.fakeHSA2AccountStatus];
 508     }
 509     [self startCKAccountStatusMock];
 510 }
 511
 512 - (void)startCKAccountStatusMock {
 513     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 514     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 515     // and return; it's just a tad confusing.
 516     if([self.ckaccountHoldOperation isPending]) {
 517         [self.operationQueue addOperation: self.ckaccountHoldOperation];
 518     }
 519
 520     [self.accountStateTracker performInitialDispatches];
 521 }
 522
 523 -(void)holdCloudKitModifications {
 524     XCTAssertFalse([self.ckModifyHoldOperation isPending], "Shouldn't already be a pending cloudkit modify hold operation");
 525     self.ckModifyHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 526         secnotice("ckks", "Released CloudKit modification hold.");
 527     }];
 528 }
 529 -(void)releaseCloudKitModificationHold {
 530     if([self.ckModifyHoldOperation isPending]) {
 531         [self.operationQueue addOperation: self.ckModifyHoldOperation];
 532     }
 533 }
 534
 535 -(void)holdCloudKitFetches {
 536     XCTAssertFalse([self.ckFetchHoldOperation isPending], "Shouldn't already be a pending cloudkit fetch hold operation");
 537     self.ckFetchHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 538         secnotice("ckks", "Released CloudKit fetch hold.");
 539     }];
 540 }
 541 -(void)releaseCloudKitFetchHold {
 542     if([self.ckFetchHoldOperation isPending]) {
 543         [self.operationQueue addOperation: self.ckFetchHoldOperation];
 544     }
 545 }
 546
 547 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID {
 548     [self expectCKModifyItemRecords:expectedNumberOfRecords
 549            currentKeyPointerRecords:expectedCurrentKeyRecords
 550                              zoneID:zoneID
 551                           checkItem:nil];
 552 }
 553
 554 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID checkItem: (BOOL (^)(CKRecord*)) checkItem {
 555     [self expectCKModifyItemRecords:expectedNumberOfRecords
 556                      deletedRecords:0
 557            currentKeyPointerRecords:expectedCurrentKeyRecords
 558                              zoneID:zoneID
 559                           checkItem:checkItem];
 560 }
 561
 562 - (void)expectCKModifyItemRecords:(NSUInteger)expectedNumberOfModifiedRecords
 563                    deletedRecords:(NSUInteger)expectedNumberOfDeletedRecords
 564          currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 565                            zoneID:(CKRecordZoneID*)zoneID
 566                         checkItem:(BOOL (^)(CKRecord*))checkItem {
 567     // We're updating the device state type on every update, so add it in here
 568     NSMutableDictionary* expectedRecords = [@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfModifiedRecords],
 569                                               SecCKRecordCurrentKeyType: [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords],
 570                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 571                                               } mutableCopy];
 572
 573     if(SecCKKSSyncManifests()) {
 574         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 1];
 575         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 576     }
 577
 578     NSDictionary* deletedRecords = nil;
 579     if(expectedNumberOfDeletedRecords != 0) {
 580         deletedRecords = @{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfDeletedRecords]};
 581     }
 582
 583     [self expectCKModifyRecords:expectedRecords
 584         deletedRecordTypeCounts:deletedRecords
 585                          zoneID:zoneID
 586             checkModifiedRecord: ^BOOL (CKRecord* record){
 587                 if([record.recordType isEqualToString: SecCKRecordItemType] && checkItem) {
 588                     return checkItem(record);
 589                 } else {
 590                     return YES;
 591                 }
 592             }
 593            runAfterModification:nil];
 594 }
 595
 596
 597
 598 - (void)expectCKModifyKeyRecords:(NSUInteger)expectedNumberOfRecords
 599         currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 600                  tlkShareRecords:(NSUInteger)expectedTLKShareRecords
 601                           zoneID:(CKRecordZoneID*)zoneID
 602 {
 603     return [self expectCKModifyKeyRecords:expectedNumberOfRecords
 604                  currentKeyPointerRecords:expectedCurrentKeyRecords
 605                           tlkShareRecords:expectedTLKShareRecords
 606                                    zoneID:zoneID
 607                       checkModifiedRecord:nil];
 608 }
 609
 610 - (void)expectCKModifyKeyRecords:(NSUInteger)expectedNumberOfRecords
 611         currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 612                  tlkShareRecords:(NSUInteger)expectedTLKShareRecords
 613                           zoneID:(CKRecordZoneID*)zoneID
 614              checkModifiedRecord:(BOOL (^_Nullable)(CKRecord*))checkModifiedRecord
 615 {
 616     NSNumber* nkeys = [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords];
 617     NSNumber* ncurrentkeys = [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords];
 618     NSNumber* ntlkshares = [NSNumber numberWithUnsignedInteger: expectedTLKShareRecords];
 619
 620     [self expectCKModifyRecords:@{SecCKRecordIntermediateKeyType: nkeys,
 621                                   SecCKRecordCurrentKeyType: ncurrentkeys,
 622                                   SecCKRecordTLKShareType: ntlkshares,
 623                                   }
 624         deletedRecordTypeCounts:nil
 625                          zoneID:zoneID
 626             checkModifiedRecord:checkModifiedRecord
 627            runAfterModification:nil];
 628 }
 629
 630 - (void)expectCKModifyRecords:(NSDictionary<NSString*, NSNumber*>*) expectedRecordTypeCounts
 631       deletedRecordTypeCounts:(NSDictionary<NSString*, NSNumber*>*) expectedDeletedRecordTypeCounts
 632                        zoneID:(CKRecordZoneID*) zoneID
 633           checkModifiedRecord:(BOOL (^)(CKRecord*)) checkModifiedRecord
 634          runAfterModification:(void (^) (void))afterModification
 635 {
 636     __weak __typeof(self) weakSelf = self;
 637
 638     // Create an object for the block to retain and modify
 639     BoolHolder* runAlready = [[BoolHolder alloc] init];
 640
 641     secnotice("fakecloudkit", "expecting an operation matching modifications: %@ deletions: %@",
 642               expectedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 643
 644     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 645         secnotice("fakecloudkit", "Received an operation (%@), checking if it's a modification", obj);
 646         __block bool matches = false;
 647         if(runAlready.state) {
 648             secnotice("fakecloudkit", "Run already, skipping");
 649             return NO;
 650         }
 651
 652         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 653             __strong __typeof(weakSelf) strongSelf = weakSelf;
 654             XCTAssertNotNil(strongSelf, "self exists");
 655
 656             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 657             matches = true;
 658
 659             NSMutableDictionary<NSString*, NSNumber*>* modifiedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 660             NSMutableDictionary<NSString*, NSNumber*>* deletedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 661
 662             // First: check if it matches. If it does, _then_ execute the operation.
 663             // Supports single-zone atomic writes only
 664
 665             if(!op.atomic) {
 666                 // We only care about atomic operations
 667                 secnotice("fakecloudkit", "Not an atomic operation; quitting: %@", op);
 668                 return NO;
 669             }
 670
 671             FakeCKZone* zone = strongSelf.zones[zoneID];
 672             XCTAssertNotNil(zone, "Have a zone for these records");
 673
 674             __block BOOL result = YES;
 675             dispatch_sync(zone.queue, ^{
 676
 677                 for(CKRecord* record in op.recordsToSave) {
 678                     if(![record.recordID.zoneID isEqual: zoneID]) {
 679                         secnotice("fakecloudkit", "Modified record zone ID mismatch: %@ %@", zoneID, record.recordID.zoneID);
 680                         result = NO;
 681                         return;
 682                     }
 683
 684                     NSError* recordError = [zone errorFromSavingRecord: record];
 685                     if(recordError) {
 686                         secnotice("fakecloudkit", "Record zone rejected record write: %@ %@", recordError, record);
 687                         XCTFail(@"Record zone rejected record write: %@ %@", recordError, record);
 688                         result = NO;
 689                         return;
 690                     }
 691
 692                     NSNumber* currentCountNumber = modifiedRecordTypeCounts[record.recordType];
 693                     NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 694                     modifiedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 695                 }
 696
 697                 for(CKRecordID* recordID in op.recordIDsToDelete) {
 698                     if(![recordID.zoneID isEqual: zoneID]) {
 699                         matches = false;
 700                         secnotice("fakecloudkit", "Deleted record zone ID mismatch: %@ %@", zoneID, recordID.zoneID);
 701                     }
 702
 703                     // Find the object in CloudKit, and record its type
 704                     CKRecord* record = strongSelf.zones[zoneID].currentDatabase[recordID];
 705                     if(record) {
 706                         NSNumber* currentCountNumber = deletedRecordTypeCounts[record.recordType];
 707                         NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 708                         deletedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 709                     }
 710                 }
 711
 712                 NSMutableDictionary* filteredExpectedRecordTypeCounts = [expectedRecordTypeCounts mutableCopy];
 713                 for(NSString* key in filteredExpectedRecordTypeCounts.allKeys) {
 714                     if([filteredExpectedRecordTypeCounts[key] isEqual: [NSNumber numberWithInt:0]]) {
 715                         filteredExpectedRecordTypeCounts[key] = nil;
 716                     }
 717                 }
 718                 filteredExpectedRecordTypeCounts[SecCKRecordManifestType] = modifiedRecordTypeCounts[SecCKRecordManifestType];
 719                 filteredExpectedRecordTypeCounts[SecCKRecordManifestLeafType] = modifiedRecordTypeCounts[SecCKRecordManifestLeafType];
 720
 721                 // Inspect that we have exactly the same records as we expect
 722                 if(expectedRecordTypeCounts) {
 723                     matches &= !![modifiedRecordTypeCounts isEqual: filteredExpectedRecordTypeCounts];
 724                     if(!matches) {
 725                         secnotice("fakecloudkit", "Record number mismatch: %@ %@", modifiedRecordTypeCounts, filteredExpectedRecordTypeCounts);
 726                         result = NO;
 727                         return;
 728                     }
 729                 } else {
 730                     matches &= op.recordsToSave.count == 0u;
 731                     if(!matches) {
 732                         secnotice("fakecloudkit", "Record number mismatch: %@ 0", modifiedRecordTypeCounts);
 733                         result = NO;
 734                         return;
 735                     }
 736                 }
 737                 if(expectedDeletedRecordTypeCounts) {
 738                     matches &= !![deletedRecordTypeCounts  isEqual: expectedDeletedRecordTypeCounts];
 739                     if(!matches) {
 740                         secnotice("fakecloudkit", "Deleted record number mismatch: %@ %@", deletedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 741                         result = NO;
 742                         return;
 743                     }
 744                 } else {
 745                     matches &= op.recordIDsToDelete.count == 0u;
 746                     if(!matches) {
 747                         secnotice("fakecloudkit", "Deleted record number mismatch: %@ 0", deletedRecordTypeCounts);
 748                         result = NO;
 749                         return;
 750                     }
 751                 }
 752
 753                 // We have the right number of things, and their etags match. Ensure that they have the right etags
 754                 if(matches && checkModifiedRecord) {
 755                     // Clearly we have the right number of things. Call checkRecord on them...
 756                     for(CKRecord* record in op.recordsToSave) {
 757                         matches &= !!(checkModifiedRecord(record));
 758                         if(!matches) {
 759                             secnotice("fakecloudkit", "Check record reports NO: %@ 0", record);
 760                             result = NO;
 761                             return;
 762                         }
 763                     }
 764                 }
 765
 766                 if(matches) {
 767                     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 768                     // if you'd like to read the data from this write.
 769                     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-write" withBlock: ^{
 770                         @synchronized(zone.currentDatabase) {
 771                             NSMutableArray* savedRecords = [[NSMutableArray alloc] init];
 772                             for(CKRecord* record in op.recordsToSave) {
 773                                 CKRecord* reflectedRecord = [record copy];
 774                                 reflectedRecord.modificationDate = [NSDate date];
 775
 776                                 [zone addToZone: reflectedRecord];
 777
 778                                 [savedRecords addObject:reflectedRecord];
 779                                 op.perRecordCompletionBlock(reflectedRecord, nil);
 780                             }
 781                             for(CKRecordID* recordID in op.recordIDsToDelete) {
 782                                 // I don't believe CloudKit fails an operation if you delete a record that's not there, so:
 783                                 [zone deleteCKRecordIDFromZone: recordID];
 784                             }
 785
 786                             if(afterModification) {
 787                                 afterModification();
 788                             }
 789
 790                             op.modifyRecordsCompletionBlock(savedRecords, op.recordIDsToDelete, nil);
 791                             op.isFinished = YES;
 792                         }
 793                     }];
 794                     [ckop addNullableDependency:strongSelf.ckModifyHoldOperation];
 795                     [strongSelf.operationQueue addOperation: ckop];
 796                 }
 797             });
 798             if(result != YES) {
 799                 return result;
 800             }
 801         }
 802         if(matches) {
 803             runAlready.state = true;
 804         }
 805         return matches ? YES : NO;
 806     }]];
 807 }
 808
 809 - (void)failNextZoneCreation:(CKRecordZoneID*)zoneID {
 810     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 811     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 812     self.zones[zoneID].creationError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain
 813                                                                         code:CKErrorNetworkUnavailable
 814                                                                     userInfo:@{
 815                                                                                CKErrorRetryAfterKey: @(0.5),
 816                                                                                }];
 817 }
 818
 819 // Report success, but don't actually create the zone.
 820 // This way, you can find ZoneNotFound errors later on
 821 - (void)failNextZoneCreationSilently:(CKRecordZoneID*)zoneID {
 822     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 823     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 824     self.zones[zoneID].failCreationSilently = true;
 825 }
 826
 827 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID {
 828     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 829     self.zones[zoneID].subscriptionError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 830 }
 831
 832 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID withError:(NSError*)error {
 833     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 834     self.zones[zoneID].subscriptionError = error;
 835 }
 836
 837 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 838     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:nil];
 839 }
 840
 841 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)(void))blockAfterReject {
 842     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:blockAfterReject withError:nil];
 843 }
 844
 845 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)(void))blockAfterReject withError:(NSError*)error {
 846     __weak __typeof(self) weakSelf = self;
 847
 848     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 849         __strong __typeof(weakSelf) strongSelf = weakSelf;
 850         XCTAssertNotNil(strongSelf, "self exists");
 851
 852         __block bool rejected = false;
 853         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 854             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 855
 856             if(!op.atomic) {
 857                 // We only care about atomic operations
 858                 return NO;
 859             }
 860
 861             // We want to only match zone updates pertaining to this zone
 862             for(CKRecord* record in op.recordsToSave) {
 863                 if(![record.recordID.zoneID isEqual: zoneID]) {
 864                     return NO;
 865                 }
 866             }
 867
 868             FakeCKZone* zone = strongSelf.zones[zoneID];
 869             XCTAssertNotNil(zone, "Have a zone for these records");
 870
 871             rejected = true;
 872
 873             if(error) {
 874                 [strongSelf rejectWrite: op withError:error];
 875             } else {
 876                 NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 877                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 878             }
 879
 880             if(blockAfterReject) {
 881                 blockAfterReject();
 882             }
 883         }
 884         return rejected ? YES : NO;
 885     }]];
 886 }
 887
 888 - (void)expectCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 889     __weak __typeof(self) weakSelf = self;
 890
 891     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 892         __strong __typeof(weakSelf) strongSelf = weakSelf;
 893         XCTAssertNotNil(strongSelf, "self exists");
 894
 895         __block bool rejected = false;
 896         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 897             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 898
 899             secnotice("fakecloudkit", "checking for expectCKAtomicModifyItemRecordsUpdateFailure");
 900
 901             if(!op.atomic) {
 902                 // We only care about atomic operations
 903                 secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: update not atomic");
 904                 return NO;
 905             }
 906
 907             // We want to only match zone updates pertaining to this zone
 908             for(CKRecord* record in op.recordsToSave) {
 909                 if(![record.recordID.zoneID isEqual: zoneID]) {
 910                     secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: %@ is not %@", record.recordID.zoneID, zoneID);
 911                     return NO;
 912                 }
 913             }
 914
 915             FakeCKZone* zone = strongSelf.zones[zoneID];
 916             XCTAssertNotNil(zone, "Have a zone for these records");
 917
 918             NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 919
 920             @synchronized(zone.currentDatabase) {
 921                 for(CKRecord* record in op.recordsToSave) {
 922                     // Check if we should allow this transaction
 923                     NSError* recordSaveError = [zone errorFromSavingRecord: record];
 924                     if(recordSaveError) {
 925                         failedRecords[record.recordID] = recordSaveError;
 926                         rejected = true;
 927                     }
 928                 }
 929             }
 930
 931             if(rejected) {
 932                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 933             } else {
 934                 secnotice("fakecloudkit", "expectCKAtomicModifyItemRecordsUpdateFailure: doesn't seem like an error to us");
 935             }
 936         }
 937         return rejected ? YES : NO;
 938     }]];
 939 }
 940
 941 -(void)rejectWrite:(CKModifyRecordsOperation*)op withError:(NSError*)error {
 942     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 943     // if you'd like to read the data from this write.
 944     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write-error" withBlock: ^{
 945         op.modifyRecordsCompletionBlock(nil, nil, error);
 946         op.isFinished = YES;
 947     }];
 948     [ckop addNullableDependency: self.ckModifyHoldOperation];
 949     [self.operationQueue addOperation: ckop];
 950 }
 951
 952 -(void)rejectWrite:(CKModifyRecordsOperation*)op failedRecords:(NSMutableDictionary<CKRecordID*, NSError*>*)failedRecords {
 953     // Add the batch request failed errors
 954     for(CKRecord* record in op.recordsToSave) {
 955         NSError* exists = failedRecords[record.recordID];
 956         if(!exists) {
 957             // TODO: might have important userInfo, but we're not mocking that yet
 958             failedRecords[record.recordID] = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorBatchRequestFailed userInfo: @{}];
 959         }
 960     }
 961
 962     NSError* error = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorPartialFailure userInfo: @{CKPartialErrorsByItemIDKey: failedRecords}];
 963
 964     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 965     // if you'd like to read the data from this write.
 966     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write" withBlock: ^{
 967         op.modifyRecordsCompletionBlock(nil, nil, error);
 968         op.isFinished = YES;
 969     }];
 970     [ckop addNullableDependency: self.ckModifyHoldOperation];
 971     [self.operationQueue addOperation: ckop];
 972 }
 973
 974 - (void)expectCKDeleteItemRecords:(NSUInteger)expectedNumberOfRecords
 975                            zoneID:(CKRecordZoneID*) zoneID {
 976
 977     // We're updating the device state type on every update, so add it in here
 978     NSMutableDictionary* expectedRecords = [@{
 979                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInteger:expectedNumberOfRecords],
 980                                               } mutableCopy];
 981     if(SecCKKSSyncManifests()) {
 982         // TODO: this really shouldn't be 2.
 983         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 2];
 984         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 985     }
 986
 987     [self expectCKModifyRecords:expectedRecords
 988         deletedRecordTypeCounts:@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords]}
 989                          zoneID:zoneID
 990             checkModifiedRecord:nil
 991            runAfterModification:nil];
 992 }
 993
 994 -(void)waitForCKModifications {
 995     // CloudKit modifications are put on the local queue.
 996     // This is heavyweight but should suffice.
 997     [self.operationQueue waitUntilAllOperationsAreFinished];
 998 }
 999
1000 - (void)tearDown {
1001     NSString* testName = [self.name componentsSeparatedByString:@" "][1];
1002     testName = [testName stringByReplacingOccurrencesOfString:@"]" withString:@""];
1003     secnotice("ckkstest", "Ending test %@", testName);
1004
1005     if(SecCKKSIsEnabled()) {
1006         self.accountStatus = CKAccountStatusCouldNotDetermine;
1007
1008         // If the test never initialized the account state, don't call status later
1009         bool callStatus = [self.ckaccountHoldOperation isFinished];
1010         [self.ckaccountHoldOperation cancel];
1011         self.ckaccountHoldOperation = nil;
1012
1013         // Ensure we don't have any blocking operations left
1014         [self.operationQueue cancelAllOperations];
1015         [self waitForCKModifications];
1016
1017         XCTAssertEqual(0, [self.injectedManager.completedSecCKKSInitialize wait:20*NSEC_PER_SEC],
1018             "Timeout did not occur waiting for SecCKKSInitialize");
1019
1020         // Ensure that we can fetch zone status for all zones
1021         if(callStatus) {
1022             XCTestExpectation *statusReturned = [self expectationWithDescription:@"status returned"];
1023             [self.injectedManager rpcStatus:nil reply:^(NSArray<NSDictionary *> *result, NSError *error) {
1024                 XCTAssertNil(error, "Should be no error fetching status");
1025                 [statusReturned fulfill];
1026             }];
1027             [self waitForExpectations: @[statusReturned] timeout:20];
1028
1029             // Make sure this happens before teardown.
1030             XCTAssertEqual(0, [self.accountStateTracker.finishedInitialDispatches wait:20*NSEC_PER_SEC], "Account state tracker initialized itself");
1031
1032             dispatch_group_t accountChangesDelivered = [self.accountStateTracker checkForAllDeliveries];
1033             XCTAssertEqual(0, dispatch_group_wait(accountChangesDelivered, dispatch_time(DISPATCH_TIME_NOW, 10*NSEC_PER_SEC)), "Account state tracker finished delivering everything");
1034         }
1035     }
1036
1037     [super tearDown];
1038
1039     [self.injectedManager cancelPendingOperations];
1040     [CKKSViewManager resetManager:true setTo:nil];
1041     self.injectedManager = nil;
1042     [self.mockCKKSViewManager stopMocking];
1043     self.mockCKKSViewManager = nil;
1044
1045     [self.mockAccountStateTracker stopMocking];
1046     self.mockAccountStateTracker = nil;
1047
1048     [self.mockLockStateTracker stopMocking];
1049     self.mockLockStateTracker = nil;
1050
1051     [self.mockFakeCKModifyRecordZonesOperation stopMocking];
1052     self.mockFakeCKModifyRecordZonesOperation = nil;
1053
1054     [self.mockFakeCKModifySubscriptionsOperation stopMocking];
1055     self.mockFakeCKModifySubscriptionsOperation = nil;
1056
1057     [self.mockFakeCKFetchRecordZoneChangesOperation stopMocking];
1058     self.mockFakeCKFetchRecordZoneChangesOperation = nil;
1059
1060     [self.mockFakeCKFetchRecordsOperation stopMocking];
1061     self.mockFakeCKFetchRecordsOperation = nil;
1062
1063     [self.mockFakeCKQueryOperation stopMocking];
1064     self.mockFakeCKQueryOperation = nil;
1065
1066     [self.mockDatabase stopMocking];
1067     self.mockDatabase = nil;
1068
1069     [self.mockDatabaseExceptionCatcher stopMocking];
1070     self.mockDatabaseExceptionCatcher = nil;
1071
1072     [self.mockContainer stopMocking];
1073     self.mockContainer = nil;
1074
1075     [self.mockTTR stopMocking];
1076     self.mockTTR = nil;
1077     self.ttrExpectation = nil;
1078     self.isTTRRatelimited = true;
1079
1080     self.zones = nil;
1081
1082     _mockSOSAdapter = nil;
1083     _mockOctagonAdapter = nil;
1084
1085     SecCKKSTestResetFlags();
1086 }
1087
1088 - (CKKSKey*) fakeTLK: (CKRecordZoneID*)zoneID {
1089     CKKSKey* key = [[CKKSKey alloc] initSelfWrappedWithAESKey:[[CKKSAESSIVKey alloc] initWithBase64: @"uImdbZ7Zg+6WJXScTnRBfNmoU1UiMkSYxWc+d1Vuq3IFn2RmTRkTdWTe3HmeWo1pAomqy+upK8KHg2PGiRGhqg=="]
1090                                                          uuid:[[NSUUID UUID] UUIDString]
1091                                                      keyclass:SecCKKSKeyClassTLK
1092                                                         state: SecCKKSProcessedStateLocal
1093                                                        zoneID:zoneID
1094                                               encodedCKRecord: nil
1095                                                    currentkey: true];
1096     [key CKRecordWithZoneID: zoneID];
1097     return key;
1098 }
1099
1100 - (NSError*)ckInternalServerExtensionError:(NSInteger)code description:(NSString*)desc {
1101     NSError* extensionError = [[CKPrettyError alloc] initWithDomain:@"CloudkitKeychainService"
1102                                                                code:code
1103                                                            userInfo:@{
1104                                                                       CKErrorServerDescriptionKey: desc,
1105                                                                       NSLocalizedDescriptionKey: desc,
1106                                                                       }];
1107     NSError* internalError = [[CKPrettyError alloc] initWithDomain:CKInternalErrorDomain
1108                                                               code:CKErrorInternalPluginError
1109                                                           userInfo:@{CKErrorServerDescriptionKey: desc,
1110                                                                      NSLocalizedDescriptionKey: desc,
1111                                                                      NSUnderlyingErrorKey: extensionError,
1112                                                                      }];
1113     NSError* error = [[CKPrettyError alloc] initWithDomain:CKErrorDomain
1114                                                       code:CKErrorServerRejectedRequest
1115                                                   userInfo:@{NSUnderlyingErrorKey: internalError,
1116                                                              CKErrorServerDescriptionKey: desc,
1117                                                              NSLocalizedDescriptionKey: desc,
1118                                                              CKContainerIDKey: SecCKKSContainerName,
1119                                                              }];
1120     return error;
1121 }
1122
1123 @end
1124
1125 #endif