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 "OSX/sec/securityd/Regressions/SecdTestKeychainUtilities.h"
  36 #include <utilities/SecFileLocations.h>
  37 #include <securityd/SecItemServer.h>
  38
  39 #if NO_SERVER
  40 #include <securityd/spi.h>
  41 #endif
  42
  43 #include <Security/SecureObjectSync/SOSViews.h>
  44
  45 #include <utilities/SecDb.h>
  46 #include <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
  56 #import "MockCloudKit.h"
  57
  58 @interface BoolHolder : NSObject
  59 @property bool state;
  60 @end
  61
  62 @implementation BoolHolder
  63 @end
  64
  65 // Inform OCMock about the internals of CKContainer
  66 @interface CKContainer ()
  67 - (void)_checkSelfCloudServicesEntitlement;
  68 @end
  69
  70
  71 @implementation CloudKitMockXCTest
  72
  73 + (void)setUp {
  74     // Turn on testing
  75     SecCKKSTestsEnable();
  76     [super setUp];
  77
  78 #if NO_SERVER
  79     securityd_init_local_spi();
  80 #endif
  81 }
  82
  83 - (void)setUp {
  84     [super setUp];
  85
  86     // All tests start with the same flag set.
  87     SecCKKSTestResetFlags();
  88     SecCKKSTestSetDisableSOS(true);
  89
  90     self.silentFetchesAllowed = true;
  91
  92     __weak __typeof(self) weakSelf = self;
  93     self.operationQueue = [[NSOperationQueue alloc] init];
  94     self.operationQueue.maxConcurrentOperationCount = 1;
  95
  96     self.zones = [[NSMutableDictionary alloc] init];
  97
  98     self.mockDatabase = OCMStrictClassMock([CKDatabase class]);
  99     self.mockContainer = OCMClassMock([CKContainer class]);
 100     OCMStub([self.mockContainer containerWithIdentifier:[OCMArg isKindOfClass:[NSString class]]]).andReturn(self.mockContainer);
 101     OCMStub([self.mockContainer defaultContainer]).andReturn(self.mockContainer);
 102     OCMStub([self.mockContainer alloc]).andReturn(self.mockContainer);
 103     OCMStub([self.mockContainer containerIdentifier]).andReturn(SecCKKSContainerName);
 104     OCMStub([self.mockContainer initWithContainerID: [OCMArg any] options: [OCMArg any]]).andReturn(self.mockContainer);
 105     OCMStub([self.mockContainer privateCloudDatabase]).andReturn(self.mockDatabase);
 106     OCMStub([self.mockContainer serverPreferredPushEnvironmentWithCompletionHandler: ([OCMArg invokeBlockWithArgs:@"fake APS push string", [NSNull null], nil])]);
 107
 108     // If you want to change this, you'll need to update the mock
 109     _ckDeviceID = @"fake-cloudkit-device-id";
 110     OCMStub([self.mockContainer fetchCurrentDeviceIDWithCompletionHandler: ([OCMArg invokeBlockWithArgs:self.ckDeviceID, [NSNull null], nil])]);
 111
 112     self.accountStatus = CKAccountStatusAvailable;
 113     self.supportsDeviceToDeviceEncryption = YES;
 114
 115     // Inject a fake operation dependency into the manager object, so that the tests can perform setup and mock expectations before zone setup begins
 116     // Also blocks all CK account state retrieval operations (but not circle status ones)
 117     self.ckaccountHoldOperation = [NSBlockOperation named:@"ckaccount-hold" withBlock:^{
 118         secnotice("ckks", "CKKS CK account status test hold released");
 119     }];
 120
 121     OCMStub([self.mockContainer accountStatusWithCompletionHandler:
 122                     [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (CKAccountStatus accountStatus,
 123                                                                       NSError * _Nullable error)) {
 124
 125         if(passedBlock) {
 126             __strong __typeof(self) strongSelf = weakSelf;
 127             NSBlockOperation* fulfillBlock = [NSBlockOperation named:@"account-status-completion" withBlock: ^{
 128                 passedBlock(weakSelf.accountStatus, nil);
 129             }];
 130             [fulfillBlock addDependency: strongSelf.ckaccountHoldOperation];
 131             [strongSelf.operationQueue addOperation: fulfillBlock];
 132
 133             return YES;
 134         }
 135         return NO;
 136     }]]);
 137
 138     OCMStub([self.mockContainer accountInfoWithCompletionHandler:
 139              [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (CKAccountInfo* accountInfo,
 140                                                                NSError * error)) {
 141         __strong __typeof(self) strongSelf = weakSelf;
 142         if(passedBlock && strongSelf) {
 143             NSBlockOperation* fulfillBlock = [NSBlockOperation named:@"account-info-completion" withBlock: ^{
 144                 __strong __typeof(self) blockStrongSelf = weakSelf;
 145                 CKAccountInfo* account = [[CKAccountInfo alloc] init];
 146                 account.accountStatus = blockStrongSelf.accountStatus;
 147                 account.supportsDeviceToDeviceEncryption = blockStrongSelf.supportsDeviceToDeviceEncryption;
 148                 account.accountPartition = CKAccountPartitionTypeProduction;
 149                 passedBlock((CKAccountInfo*)account, nil);
 150             }];
 151             [fulfillBlock addDependency: strongSelf.ckaccountHoldOperation];
 152             [strongSelf.operationQueue addOperation: fulfillBlock];
 153
 154             return YES;
 155         }
 156         return NO;
 157     }]]);
 158
 159     self.circleStatus = kSOSCCInCircle;
 160     self.mockAccountStateTracker = OCMClassMock([CKKSCKAccountStateTracker class]);
 161     OCMStub([self.mockAccountStateTracker getCircleStatus]).andCall(self, @selector(circleStatus));
 162
 163     // If we're in circle, come up with a fake circle id. Otherwise, return an error.
 164     self.circlePeerID = @"fake-circle-id";
 165     OCMStub([self.mockAccountStateTracker fetchCirclePeerID:
 166              [OCMArg checkWithBlock:^BOOL(void (^passedBlock) (NSString* peerID,
 167                                                                NSError * error)) {
 168         __strong __typeof(self) strongSelf = weakSelf;
 169         if(passedBlock && strongSelf) {
 170             if(strongSelf.circleStatus == kSOSCCInCircle) {
 171                 passedBlock(strongSelf.circlePeerID, nil);
 172             } else {
 173                 passedBlock(nil, [NSError errorWithDomain:@"securityd" code:errSecInternalError userInfo:@{NSLocalizedDescriptionKey:@"no account, no circle id"}]);
 174             }
 175
 176             return YES;
 177         }
 178         return NO;
 179     }]]);
 180
 181     self.aksLockState = false; // Lie and say AKS is always unlocked
 182     self.mockLockStateTracker = OCMClassMock([CKKSLockStateTracker class]);
 183     OCMStub([self.mockLockStateTracker queryAKSLocked]).andCall(self, @selector(aksLockState));
 184
 185     self.mockFakeCKModifyRecordZonesOperation = OCMClassMock([FakeCKModifyRecordZonesOperation class]);
 186     OCMStub([self.mockFakeCKModifyRecordZonesOperation ckdb]).andReturn(self.zones);
 187
 188     self.mockFakeCKModifySubscriptionsOperation = OCMClassMock([FakeCKModifySubscriptionsOperation class]);
 189     OCMStub([self.mockFakeCKModifySubscriptionsOperation ckdb]).andReturn(self.zones);
 190
 191     self.mockFakeCKFetchRecordZoneChangesOperation = OCMClassMock([FakeCKFetchRecordZoneChangesOperation class]);
 192     OCMStub([self.mockFakeCKFetchRecordZoneChangesOperation ckdb]).andReturn(self.zones);
 193
 194     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 195         __strong __typeof(self) strongSelf = weakSelf;
 196         BOOL matches = NO;
 197         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 198             if(strongSelf.silentFetchesAllowed) {
 199                 matches = YES;
 200
 201                 FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 202                 [strongSelf.operationQueue addOperation: frzco];
 203             }
 204         }
 205         return matches;
 206     }]]);
 207
 208
 209     self.testZoneID = [[CKRecordZoneID alloc] initWithZoneName:@"testzone" ownerName:CKCurrentUserDefaultName];
 210
 211     // Inject a fake operation dependency into the manager object, so that the tests can perform setup and mock expectations before zone setup begins
 212     // Also blocks all CK account state retrieval operations (but not circle status ones)
 213     self.ckksHoldOperation = [[NSBlockOperation alloc] init];
 214     [self.ckksHoldOperation addExecutionBlock:^{
 215         secnotice("ckks", "CKKS testing hold released");
 216     }];
 217     self.ckksHoldOperation.name = @"ckks-hold";
 218
 219     self.mockCKKSViewManager = OCMClassMock([CKKSViewManager class]);
 220     OCMStub([self.mockCKKSViewManager viewList]).andCall(self, @selector(managedViewList));
 221     OCMStub([self.mockCKKSViewManager syncBackupAndNotifyAboutSync]);
 222
 223     self.injectedManager = [[CKKSViewManager alloc] initWithContainerName:SecCKKSContainerName
 224                                                                    usePCS:SecCKKSContainerUsePCS
 225                                      fetchRecordZoneChangesOperationClass:[FakeCKFetchRecordZoneChangesOperation class]
 226                                         modifySubscriptionsOperationClass:[FakeCKModifySubscriptionsOperation class]
 227                                           modifyRecordZonesOperationClass:[FakeCKModifyRecordZonesOperation class]
 228                                                        apsConnectionClass:[FakeAPSConnection class]
 229                                                 nsnotificationCenterClass:[FakeNSNotificationCenter class]
 230                                                             notifierClass:[FakeCKKSNotifier class]
 231                                                                 setupHold:self.ckksHoldOperation];
 232
 233     [CKKSViewManager resetManager:false setTo:self.injectedManager];
 234
 235     // Make a new fake keychain
 236     NSString* smallName = [self.name componentsSeparatedByString:@" "][1];
 237     smallName = [smallName stringByReplacingOccurrencesOfString:@"]" withString:@""];
 238
 239     NSString* tmp_dir = [NSString stringWithFormat: @"/tmp/%@.%X", smallName, arc4random()];
 240     [[NSFileManager defaultManager] createDirectoryAtPath:[NSString stringWithFormat: @"%@/Library/Keychains", tmp_dir] withIntermediateDirectories:YES attributes:nil error:NULL];
 241
 242     SetCustomHomeURLString((__bridge CFStringRef) tmp_dir);
 243     SecKeychainDbReset(NULL);
 244
 245     // Actually load the database.
 246     kc_with_dbt(true, NULL, ^bool (SecDbConnectionRef dbt) { return false; });
 247 }
 248
 249 -(CKKSCKAccountStateTracker*)accountStateTracker {
 250     return self.injectedManager.accountTracker;
 251 }
 252
 253 -(CKKSLockStateTracker*)lockStateTracker {
 254     return self.injectedManager.lockStateTracker;
 255 }
 256
 257 -(NSSet*)managedViewList {
 258     return (NSSet*) CFBridgingRelease(SOSViewCopyViewSet(kViewSetCKKS));
 259 }
 260
 261 -(void)expectCKFetch {
 262     // Create an object for the block to retain and modify
 263     BoolHolder* runAlready = [[BoolHolder alloc] init];
 264
 265     __weak __typeof(self) weakSelf = self;
 266     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 267         __strong __typeof(self) strongSelf = weakSelf;
 268         if(runAlready.state) {
 269             return NO;
 270         }
 271         BOOL matches = NO;
 272         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 273             matches = YES;
 274             runAlready.state = true;
 275
 276             FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 277             [strongSelf.operationQueue addOperation: frzco];
 278         }
 279         return matches;
 280     }]];
 281 }
 282
 283 - (void)startCKKSSubsystem {
 284     [self startCKAccountStatusMock];
 285     [self startCKKSSubsystemOnly];
 286 }
 287
 288 - (void)startCKKSSubsystemOnly {
 289     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 290     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 291     // and return; it's just a tad confusing.
 292     if([self.ckksHoldOperation isPending]) {
 293         [self.operationQueue addOperation: self.ckksHoldOperation];
 294     }
 295 }
 296
 297 - (void)startCKAccountStatusMock {
 298     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 299     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 300     // and return; it's just a tad confusing.
 301     if([self.ckaccountHoldOperation isPending]) {
 302         [self.operationQueue addOperation: self.ckaccountHoldOperation];
 303     }
 304 }
 305
 306 -(void)holdCloudKitModifications {
 307     self.ckModifyHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 308         secnotice("ckks", "Released CloudKit modification hold.");
 309     }];
 310 }
 311 -(void)releaseCloudKitModificationHold {
 312     if([self.ckModifyHoldOperation isPending]) {
 313         [self.operationQueue addOperation: self.ckModifyHoldOperation];
 314     }
 315 }
 316
 317 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID {
 318     [self expectCKModifyItemRecords:expectedNumberOfRecords
 319            currentKeyPointerRecords:expectedCurrentKeyRecords
 320                              zoneID:zoneID
 321                           checkItem:nil];
 322 }
 323
 324 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID checkItem: (BOOL (^)(CKRecord*)) checkItem {
 325     [self expectCKModifyItemRecords:expectedNumberOfRecords
 326                      deletedRecords:0
 327            currentKeyPointerRecords:expectedCurrentKeyRecords
 328                              zoneID:zoneID
 329                           checkItem:checkItem];
 330 }
 331
 332 - (void)expectCKModifyItemRecords:(NSUInteger)expectedNumberOfModifiedRecords
 333                    deletedRecords:(NSUInteger)expectedNumberOfDeletedRecords
 334          currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 335                            zoneID:(CKRecordZoneID*)zoneID
 336                         checkItem:(BOOL (^)(CKRecord*))checkItem {
 337     // We're updating the device state type on every update, so add it in here
 338     NSMutableDictionary* expectedRecords = [@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfModifiedRecords],
 339                                               SecCKRecordCurrentKeyType: [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords],
 340                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 341                                               } mutableCopy];
 342
 343     if(SecCKKSSyncManifests()) {
 344         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 1];
 345         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 346     }
 347
 348     NSDictionary* deletedRecords = nil;
 349     if(expectedNumberOfDeletedRecords != 0) {
 350         deletedRecords = @{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfDeletedRecords]};
 351     }
 352
 353     [self expectCKModifyRecords:expectedRecords
 354         deletedRecordTypeCounts:deletedRecords
 355                          zoneID:zoneID
 356             checkModifiedRecord: ^BOOL (CKRecord* record){
 357                 if([record.recordType isEqualToString: SecCKRecordItemType] && checkItem) {
 358                     return checkItem(record);
 359                 } else {
 360                     return YES;
 361                 }
 362             }
 363            runAfterModification:nil];
 364 }
 365
 366
 367
 368 - (void)expectCKModifyKeyRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID {
 369     NSNumber* nkeys = [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords];
 370     NSNumber* ncurrentkeys = [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords];
 371
 372     [self expectCKModifyRecords:@{SecCKRecordIntermediateKeyType: nkeys, SecCKRecordCurrentKeyType: ncurrentkeys}
 373         deletedRecordTypeCounts:nil
 374                          zoneID:zoneID
 375             checkModifiedRecord:nil
 376            runAfterModification:nil];
 377 }
 378
 379 - (void)expectCKModifyRecords:(NSDictionary<NSString*, NSNumber*>*) expectedRecordTypeCounts
 380       deletedRecordTypeCounts:(NSDictionary<NSString*, NSNumber*>*) expectedDeletedRecordTypeCounts
 381                        zoneID:(CKRecordZoneID*) zoneID
 382           checkModifiedRecord:(BOOL (^)(CKRecord*)) checkModifiedRecord
 383          runAfterModification:(void (^) ())afterModification
 384 {
 385     __weak __typeof(self) weakSelf = self;
 386
 387     // Create an object for the block to retain and modify
 388     BoolHolder* runAlready = [[BoolHolder alloc] init];
 389
 390     secnotice("fakecloudkit", "expecting an operation matching modifications: %@ deletions: %@",
 391               expectedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 392
 393     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 394         secnotice("fakecloudkit", "Received an operation, checking");
 395         __block bool matches = false;
 396         if(runAlready.state) {
 397             secnotice("fakecloudkit", "Run already, skipping");
 398             return NO;
 399         }
 400
 401         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 402             __strong __typeof(weakSelf) strongSelf = weakSelf;
 403             XCTAssertNotNil(strongSelf, "self exists");
 404
 405             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 406             matches = true;
 407
 408             NSMutableDictionary<NSString*, NSNumber*>* modifiedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 409             NSMutableDictionary<NSString*, NSNumber*>* deletedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 410
 411             // First: check if it matches. If it does, _then_ execute the operation.
 412             // Supports single-zone atomic writes only
 413
 414             if(!op.atomic) {
 415                 // We only care about atomic operations
 416                 secnotice("fakecloudkit", "Not an atomic operation; quitting: %@", op);
 417                 return NO;
 418             }
 419
 420             FakeCKZone* zone = strongSelf.zones[zoneID];
 421             XCTAssertNotNil(zone, "Have a zone for these records");
 422
 423             for(CKRecord* record in op.recordsToSave) {
 424                 if(![record.recordID.zoneID isEqual: zoneID]) {
 425                     secnotice("fakecloudkit", "Modified record zone ID mismatch: %@ %@", zoneID, record.recordID.zoneID);
 426                     return NO;
 427                 }
 428
 429                 if([zone errorFromSavingRecord: record]) {
 430                     secnotice("fakecloudkit", "Record zone rejected record write: %@", record);
 431                     return NO;
 432                 }
 433
 434                 NSNumber* currentCountNumber = modifiedRecordTypeCounts[record.recordType];
 435                 NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 436                 modifiedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 437             }
 438
 439             for(CKRecordID* recordID in op.recordIDsToDelete) {
 440                 if(![recordID.zoneID isEqual: zoneID]) {
 441                     matches = false;
 442                     secnotice("fakecloudkit", "Deleted record zone ID mismatch: %@ %@", zoneID, recordID.zoneID);
 443                 }
 444
 445                 // Find the object in CloudKit, and record its type
 446                 CKRecord* record = strongSelf.zones[zoneID].currentDatabase[recordID];
 447                 if(record) {
 448                     NSNumber* currentCountNumber = deletedRecordTypeCounts[record.recordType];
 449                     NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 450                     deletedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 451                 }
 452             }
 453
 454             NSMutableDictionary* filteredExpectedRecordTypeCounts = [expectedRecordTypeCounts mutableCopy];
 455             for(NSString* key in filteredExpectedRecordTypeCounts.allKeys) {
 456                 if([filteredExpectedRecordTypeCounts[key] isEqual: [NSNumber numberWithInt:0]]) {
 457                     filteredExpectedRecordTypeCounts[key] = nil;
 458                 }
 459             }
 460             filteredExpectedRecordTypeCounts[SecCKRecordManifestType] = modifiedRecordTypeCounts[SecCKRecordManifestType];
 461             filteredExpectedRecordTypeCounts[SecCKRecordManifestLeafType] = modifiedRecordTypeCounts[SecCKRecordManifestLeafType];
 462
 463             // Inspect that we have exactly the same records as we expect
 464             if(expectedRecordTypeCounts) {
 465                 matches &= !![modifiedRecordTypeCounts isEqual: filteredExpectedRecordTypeCounts];
 466                 if(!matches) {
 467                     secnotice("fakecloudkit", "Record number mismatch: %@ %@", modifiedRecordTypeCounts, filteredExpectedRecordTypeCounts);
 468                     return NO;
 469                 }
 470             } else {
 471                 matches &= op.recordsToSave.count == 0u;
 472                 if(!matches) {
 473                     secnotice("fakecloudkit", "Record number mismatch: %@ 0", modifiedRecordTypeCounts);
 474                     return NO;
 475                 }
 476             }
 477             if(expectedDeletedRecordTypeCounts) {
 478                 matches &= !![deletedRecordTypeCounts  isEqual: expectedDeletedRecordTypeCounts];
 479                 if(!matches) {
 480                     secnotice("fakecloudkit", "Deleted record number mismatch: %@ %@", deletedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 481                     return NO;
 482                 }
 483             } else {
 484                 matches &= op.recordIDsToDelete.count == 0u;
 485                 if(!matches) {
 486                     secnotice("fakecloudkit", "Deleted record number mismatch: %@ 0", deletedRecordTypeCounts);
 487                     return NO;
 488                 }
 489             }
 490
 491             // We have the right number of things, and their etags match. Ensure that they have the right etags
 492             if(matches && checkModifiedRecord) {
 493                 // Clearly we have the right number of things. Call checkRecord on them...
 494                 for(CKRecord* record in op.recordsToSave) {
 495                     matches &= !!(checkModifiedRecord(record));
 496                     if(!matches) {
 497                         secnotice("fakecloudkit", "Check record reports NO: %@ 0", record);
 498                         return NO;
 499                     }
 500                 }
 501             }
 502
 503             if(matches) {
 504                 // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 505                 // if you'd like to read the data from this write.
 506                 NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-write" withBlock: ^{
 507                     @synchronized(zone.currentDatabase) {
 508                         NSMutableArray* savedRecords = [[NSMutableArray alloc] init];
 509                         for(CKRecord* record in op.recordsToSave) {
 510                             CKRecord* reflectedRecord = [record copy];
 511                             reflectedRecord.modificationDate = [NSDate date];
 512
 513                             [zone addToZone: reflectedRecord];
 514
 515                             [savedRecords addObject:reflectedRecord];
 516                             op.perRecordCompletionBlock(reflectedRecord, nil);
 517                         }
 518                         for(CKRecordID* recordID in op.recordIDsToDelete) {
 519                             // I don't believe CloudKit fails an operation if you delete a record that's not there, so:
 520                             [zone deleteCKRecordIDFromZone: recordID];
 521                         }
 522
 523                         op.modifyRecordsCompletionBlock(savedRecords, op.recordIDsToDelete, nil);
 524
 525                         if(afterModification) {
 526                             afterModification();
 527                         }
 528
 529                         op.isFinished = YES;
 530                     }
 531                 }];
 532                 [ckop addNullableDependency:strongSelf.ckModifyHoldOperation];
 533                 [strongSelf.operationQueue addOperation: ckop];
 534             }
 535         }
 536         if(matches) {
 537             runAlready.state = true;
 538         }
 539         return matches ? YES : NO;
 540     }]];
 541 }
 542
 543 - (void)failNextZoneCreation:(CKRecordZoneID*)zoneID {
 544     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 545     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 546     self.zones[zoneID].creationError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 547 }
 548
 549 // Report success, but don't actually create the zone.
 550 // This way, you can find ZoneNotFound errors later on
 551 - (void)failNextZoneCreationSilently:(CKRecordZoneID*)zoneID {
 552     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 553     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 554     self.zones[zoneID].failCreationSilently = true;
 555 }
 556
 557 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID {
 558     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 559     self.zones[zoneID].subscriptionError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 560 }
 561
 562 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID withError:(NSError*)error {
 563     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 564     self.zones[zoneID].subscriptionError = error;
 565 }
 566
 567 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 568     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:nil];
 569 }
 570
 571 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)())blockAfterReject {
 572     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:blockAfterReject withError:nil];
 573 }
 574
 575 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)())blockAfterReject withError:(NSError*)error {
 576     __weak __typeof(self) weakSelf = self;
 577
 578     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 579         __strong __typeof(weakSelf) strongSelf = weakSelf;
 580         XCTAssertNotNil(strongSelf, "self exists");
 581
 582         __block bool rejected = false;
 583         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 584             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 585
 586             if(!op.atomic) {
 587                 // We only care about atomic operations
 588                 return NO;
 589             }
 590
 591             // We want to only match zone updates pertaining to this zone
 592             for(CKRecord* record in op.recordsToSave) {
 593                 if(![record.recordID.zoneID isEqual: zoneID]) {
 594                     return NO;
 595                 }
 596             }
 597
 598             FakeCKZone* zone = strongSelf.zones[zoneID];
 599             XCTAssertNotNil(zone, "Have a zone for these records");
 600
 601             rejected = true;
 602
 603             if(error) {
 604                 [strongSelf rejectWrite: op withError:error];
 605             } else {
 606                 NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 607                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 608             }
 609
 610             if(blockAfterReject) {
 611                 blockAfterReject();
 612             }
 613         }
 614         return rejected ? YES : NO;
 615     }]];
 616 }
 617
 618 - (void)expectCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 619     __weak __typeof(self) weakSelf = self;
 620
 621     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 622         __strong __typeof(weakSelf) strongSelf = weakSelf;
 623         XCTAssertNotNil(strongSelf, "self exists");
 624
 625         __block bool rejected = false;
 626         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 627             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 628
 629             if(!op.atomic) {
 630                 // We only care about atomic operations
 631                 return NO;
 632             }
 633
 634             // We want to only match zone updates pertaining to this zone
 635             for(CKRecord* record in op.recordsToSave) {
 636                 if(![record.recordID.zoneID isEqual: zoneID]) {
 637                     return NO;
 638                 }
 639             }
 640
 641             FakeCKZone* zone = strongSelf.zones[zoneID];
 642             XCTAssertNotNil(zone, "Have a zone for these records");
 643
 644             NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 645
 646             @synchronized(zone.currentDatabase) {
 647                 for(CKRecord* record in op.recordsToSave) {
 648                     // Check if we should allow this transaction
 649                     NSError* recordSaveError = [zone errorFromSavingRecord: record];
 650                     if(recordSaveError) {
 651                         failedRecords[record.recordID] = recordSaveError;
 652                         rejected = true;
 653                     }
 654                 }
 655             }
 656
 657             if(rejected) {
 658                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 659             }
 660         }
 661         return rejected ? YES : NO;
 662     }]];
 663 }
 664
 665 -(void)rejectWrite:(CKModifyRecordsOperation*)op withError:(NSError*)error {
 666     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 667     // if you'd like to read the data from this write.
 668     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write-error" withBlock: ^{
 669         op.modifyRecordsCompletionBlock(nil, nil, error);
 670         op.isFinished = YES;
 671     }];
 672     [ckop addNullableDependency: self.ckModifyHoldOperation];
 673     [self.operationQueue addOperation: ckop];
 674 }
 675
 676 -(void)rejectWrite:(CKModifyRecordsOperation*)op failedRecords:(NSMutableDictionary<CKRecordID*, NSError*>*)failedRecords {
 677     // Add the batch request failed errors
 678     for(CKRecord* record in op.recordsToSave) {
 679         NSError* exists = failedRecords[record.recordID];
 680         if(!exists) {
 681             // TODO: might have important userInfo, but we're not mocking that yet
 682             failedRecords[record.recordID] = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorBatchRequestFailed userInfo: @{}];
 683         }
 684     }
 685
 686     NSError* error = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorPartialFailure userInfo: @{CKPartialErrorsByItemIDKey: failedRecords}];
 687
 688     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 689     // if you'd like to read the data from this write.
 690     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write" withBlock: ^{
 691         op.modifyRecordsCompletionBlock(nil, nil, error);
 692         op.isFinished = YES;
 693     }];
 694     [ckop addNullableDependency: self.ckModifyHoldOperation];
 695     [self.operationQueue addOperation: ckop];
 696 }
 697
 698 - (void)expectCKDeleteItemRecords:(NSUInteger)expectedNumberOfRecords
 699                            zoneID:(CKRecordZoneID*) zoneID {
 700
 701     // We're updating the device state type on every update, so add it in here
 702     NSMutableDictionary* expectedRecords = [@{
 703                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 704                                               } mutableCopy];
 705     if(SecCKKSSyncManifests()) {
 706         // TODO: this really shouldn't be 2.
 707         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 2];
 708         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 709     }
 710
 711     [self expectCKModifyRecords:expectedRecords
 712         deletedRecordTypeCounts:@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords]}
 713                          zoneID:zoneID
 714             checkModifiedRecord:nil
 715            runAfterModification:nil];
 716 }
 717
 718 -(void)waitForCKModifications {
 719     // CloudKit modifications are put on the local queue.
 720     // This is heavyweight but should suffice.
 721     [self.operationQueue waitUntilAllOperationsAreFinished];
 722 }
 723
 724 - (void)tearDown {
 725     // Put teardown code here. This method is called after the invocation of each test method in the class.
 726
 727     if(SecCKKSIsEnabled()) {
 728         // Ensure we don't have any blocking operations
 729         [self startCKKSSubsystem];
 730
 731         [self waitForCKModifications];
 732
 733         XCTAssertEqual(0, [self.injectedManager.completedSecCKKSInitialize wait:2*NSEC_PER_SEC],
 734             "Timeout did not occur waiting for SecCKKSInitialize");
 735
 736         // Make sure this happens before teardown.
 737         XCTAssertEqual(0, [self.accountStateTracker.finishedInitialCalls wait:1*NSEC_PER_SEC], "Account state tracker initialized itself");
 738     }
 739
 740     [super tearDown];
 741
 742     [self.injectedManager cancelPendingOperations];
 743     [CKKSViewManager resetManager:true setTo:nil];
 744     self.injectedManager = nil;
 745
 746     [self.mockAccountStateTracker stopMocking];
 747     self.mockAccountStateTracker = nil;
 748
 749     [self.mockLockStateTracker stopMocking];
 750     self.mockLockStateTracker = nil;
 751
 752     [self.mockCKKSViewManager stopMocking];
 753     self.mockCKKSViewManager = nil;
 754
 755     [self.mockFakeCKModifyRecordZonesOperation stopMocking];
 756     self.mockFakeCKModifyRecordZonesOperation = nil;
 757
 758     [self.mockFakeCKModifySubscriptionsOperation stopMocking];
 759     self.mockFakeCKModifySubscriptionsOperation = nil;
 760
 761     [self.mockFakeCKFetchRecordZoneChangesOperation stopMocking];
 762     self.mockFakeCKFetchRecordZoneChangesOperation = nil;
 763
 764     [self.mockDatabase stopMocking];
 765     self.mockDatabase = nil;
 766
 767     [self.mockContainer stopMocking];
 768     self.mockContainer = nil;
 769
 770     self.zones = nil;
 771     self.operationQueue = nil;
 772     self.ckksHoldOperation = nil;
 773     self.ckaccountHoldOperation = nil;
 774
 775     SecCKKSTestResetFlags();
 776 }
 777
 778 - (CKKSKey*) fakeTLK: (CKRecordZoneID*)zoneID {
 779     CKKSKey* key = [[CKKSKey alloc] initSelfWrappedWithAESKey:[[CKKSAESSIVKey alloc] initWithBase64: @"uImdbZ7Zg+6WJXScTnRBfNmoU1UiMkSYxWc+d1Vuq3IFn2RmTRkTdWTe3HmeWo1pAomqy+upK8KHg2PGiRGhqg=="]
 780                                                          uuid:[[NSUUID UUID] UUIDString]
 781                                                      keyclass:SecCKKSKeyClassTLK
 782                                                         state: SecCKKSProcessedStateLocal
 783                                                        zoneID:zoneID
 784                                               encodedCKRecord: nil
 785                                                    currentkey: true];
 786     [key CKRecordWithZoneID: zoneID];
 787     return key;
 788 }
 789
 790
 791 @end
 792
 793 #endif