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     self.mockFakeCKFetchRecordsOperation = OCMClassMock([FakeCKFetchRecordsOperation class]);
 195     OCMStub([self.mockFakeCKFetchRecordsOperation ckdb]).andReturn(self.zones);
 196
 197     self.mockFakeCKQueryOperation = OCMClassMock([FakeCKQueryOperation class]);
 198     OCMStub([self.mockFakeCKQueryOperation ckdb]).andReturn(self.zones);
 199
 200
 201     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 202         __strong __typeof(self) strongSelf = weakSelf;
 203         BOOL matches = NO;
 204         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 205             if(strongSelf.silentFetchesAllowed) {
 206                 matches = YES;
 207
 208                 FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 209                 [frzco addNullableDependency:strongSelf.ckFetchHoldOperation];
 210                 [strongSelf.operationQueue addOperation: frzco];
 211             }
 212         }
 213         return matches;
 214     }]]);
 215
 216     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 217         __strong __typeof(self) strongSelf = weakSelf;
 218         BOOL matches = NO;
 219         if ([obj isKindOfClass: [FakeCKFetchRecordsOperation class]]) {
 220             if(strongSelf.silentFetchesAllowed) {
 221                 matches = YES;
 222
 223                 FakeCKFetchRecordsOperation *ffro = (FakeCKFetchRecordsOperation *)obj;
 224                 [ffro addNullableDependency:strongSelf.ckFetchHoldOperation];
 225                 [strongSelf.operationQueue addOperation: ffro];
 226             }
 227         }
 228         return matches;
 229     }]]);
 230
 231     OCMStub([self.mockDatabase addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 232         __strong __typeof(self) strongSelf = weakSelf;
 233         BOOL matches = NO;
 234         if ([obj isKindOfClass: [FakeCKQueryOperation class]]) {
 235             if(strongSelf.silentFetchesAllowed) {
 236                 matches = YES;
 237
 238                 FakeCKQueryOperation *fqo = (FakeCKQueryOperation *)obj;
 239                 [fqo addNullableDependency:strongSelf.ckFetchHoldOperation];
 240                 [strongSelf.operationQueue addOperation: fqo];
 241             }
 242         }
 243         return matches;
 244     }]]);
 245
 246
 247     self.testZoneID = [[CKRecordZoneID alloc] initWithZoneName:@"testzone" ownerName:CKCurrentUserDefaultName];
 248
 249     // Inject a fake operation dependency into the manager object, so that the tests can perform setup and mock expectations before zone setup begins
 250     // Also blocks all CK account state retrieval operations (but not circle status ones)
 251     self.ckksHoldOperation = [[NSBlockOperation alloc] init];
 252     [self.ckksHoldOperation addExecutionBlock:^{
 253         secnotice("ckks", "CKKS testing hold released");
 254     }];
 255     self.ckksHoldOperation.name = @"ckks-hold";
 256
 257     //self.mockCKKSViewManagerClass = OCMClassMock([CKKSViewManager class]);
 258
 259     // We don't want to use class mocks here, because they don't play well with partial mocks
 260     self.mockCKKSViewManager = OCMPartialMock(
 261         [[CKKSViewManager alloc] initWithContainerName:SecCKKSContainerName
 262                                                 usePCS:SecCKKSContainerUsePCS
 263                   fetchRecordZoneChangesOperationClass:[FakeCKFetchRecordZoneChangesOperation class]
 264                             fetchRecordsOperationClass:[FakeCKFetchRecordsOperation class]
 265                                    queryOperationClass:[FakeCKQueryOperation class]
 266                      modifySubscriptionsOperationClass:[FakeCKModifySubscriptionsOperation class]
 267                        modifyRecordZonesOperationClass:[FakeCKModifyRecordZonesOperation class]
 268                                     apsConnectionClass:[FakeAPSConnection class]
 269                              nsnotificationCenterClass:[FakeNSNotificationCenter class]
 270                                          notifierClass:[FakeCKKSNotifier class]
 271                                              setupHold:self.ckksHoldOperation]);
 272
 273     OCMStub([self.mockCKKSViewManager viewList]).andCall(self, @selector(managedViewList));
 274     OCMStub([self.mockCKKSViewManager syncBackupAndNotifyAboutSync]);
 275
 276     self.injectedManager = self.mockCKKSViewManager;
 277
 278     [CKKSViewManager resetManager:false setTo:self.injectedManager];
 279
 280     // Make a new fake keychain
 281     NSString* smallName = [self.name componentsSeparatedByString:@" "][1];
 282     smallName = [smallName stringByReplacingOccurrencesOfString:@"]" withString:@""];
 283
 284     NSString* tmp_dir = [NSString stringWithFormat: @"/tmp/%@.%X", smallName, arc4random()];
 285     [[NSFileManager defaultManager] createDirectoryAtPath:[NSString stringWithFormat: @"%@/Library/Keychains", tmp_dir] withIntermediateDirectories:YES attributes:nil error:NULL];
 286
 287     SetCustomHomeURLString((__bridge CFStringRef) tmp_dir);
 288     SecKeychainDbReset(NULL);
 289
 290     // Actually load the database.
 291     kc_with_dbt(true, NULL, ^bool (SecDbConnectionRef dbt) { return false; });
 292 }
 293
 294 -(CKKSCKAccountStateTracker*)accountStateTracker {
 295     return self.injectedManager.accountTracker;
 296 }
 297
 298 -(CKKSLockStateTracker*)lockStateTracker {
 299     return self.injectedManager.lockStateTracker;
 300 }
 301
 302 -(NSSet*)managedViewList {
 303     return (NSSet*) CFBridgingRelease(SOSViewCopyViewSet(kViewSetCKKS));
 304 }
 305
 306 -(void)expectCKFetch {
 307     [self expectCKFetchAndRunBeforeFinished: nil];
 308 }
 309
 310 -(void)expectCKFetchAndRunBeforeFinished: (void (^)())blockAfterFetch {
 311     // Create an object for the block to retain and modify
 312     BoolHolder* runAlready = [[BoolHolder alloc] init];
 313
 314     __weak __typeof(self) weakSelf = self;
 315     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 316         __strong __typeof(self) strongSelf = weakSelf;
 317         if(runAlready.state) {
 318             return NO;
 319         }
 320         BOOL matches = NO;
 321         if ([obj isKindOfClass: [FakeCKFetchRecordZoneChangesOperation class]]) {
 322             matches = YES;
 323             runAlready.state = true;
 324
 325             FakeCKFetchRecordZoneChangesOperation *frzco = (FakeCKFetchRecordZoneChangesOperation *)obj;
 326             frzco.blockAfterFetch = blockAfterFetch;
 327             [frzco addNullableDependency: strongSelf.ckFetchHoldOperation];
 328             [strongSelf.operationQueue addOperation: frzco];
 329         }
 330         return matches;
 331     }]];
 332 }
 333
 334 -(void)expectCKFetchByRecordID {
 335     // Create an object for the block to retain and modify
 336     BoolHolder* runAlready = [[BoolHolder alloc] init];
 337
 338     __weak __typeof(self) weakSelf = self;
 339     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 340         __strong __typeof(self) strongSelf = weakSelf;
 341         if(runAlready.state) {
 342             return NO;
 343         }
 344         BOOL matches = NO;
 345         if ([obj isKindOfClass: [FakeCKFetchRecordsOperation class]]) {
 346             matches = YES;
 347             runAlready.state = true;
 348
 349             FakeCKFetchRecordsOperation *ffro = (FakeCKFetchRecordsOperation *)obj;
 350             [ffro addNullableDependency: strongSelf.ckFetchHoldOperation];
 351             [strongSelf.operationQueue addOperation: ffro];
 352         }
 353         return matches;
 354     }]];
 355 }
 356
 357
 358 -(void)expectCKFetchByQuery {
 359     // Create an object for the block to retain and modify
 360     BoolHolder* runAlready = [[BoolHolder alloc] init];
 361
 362     __weak __typeof(self) weakSelf = self;
 363     [[self.mockDatabase expect] addOperation: [OCMArg checkWithBlock:^BOOL(id obj) {
 364         __strong __typeof(self) strongSelf = weakSelf;
 365         if(runAlready.state) {
 366             return NO;
 367         }
 368         BOOL matches = NO;
 369         if ([obj isKindOfClass: [FakeCKQueryOperation class]]) {
 370             matches = YES;
 371             runAlready.state = true;
 372
 373             FakeCKQueryOperation *fqo = (FakeCKQueryOperation *)obj;
 374             [fqo addNullableDependency: strongSelf.ckFetchHoldOperation];
 375             [strongSelf.operationQueue addOperation: fqo];
 376         }
 377         return matches;
 378     }]];
 379 }
 380
 381 - (void)startCKKSSubsystem {
 382     [self startCKAccountStatusMock];
 383     [self startCKKSSubsystemOnly];
 384 }
 385
 386 - (void)startCKKSSubsystemOnly {
 387     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 388     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 389     // and return; it's just a tad confusing.
 390     if([self.ckksHoldOperation isPending]) {
 391         [self.operationQueue addOperation: self.ckksHoldOperation];
 392     }
 393 }
 394
 395 - (void)startCKAccountStatusMock {
 396     // Note: currently, based on how we're mocking up the zone creation and zone subscription operation,
 397     // they will 'fire' before this method is called. It's harmless, since the mocks immediately succeed
 398     // and return; it's just a tad confusing.
 399     if([self.ckaccountHoldOperation isPending]) {
 400         [self.operationQueue addOperation: self.ckaccountHoldOperation];
 401     }
 402 }
 403
 404 -(void)holdCloudKitModifications {
 405     XCTAssertFalse([self.ckModifyHoldOperation isPending], "Shouldn't already be a pending cloudkit modify hold operation");
 406     self.ckModifyHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 407         secnotice("ckks", "Released CloudKit modification hold.");
 408     }];
 409 }
 410 -(void)releaseCloudKitModificationHold {
 411     if([self.ckModifyHoldOperation isPending]) {
 412         [self.operationQueue addOperation: self.ckModifyHoldOperation];
 413     }
 414 }
 415
 416 -(void)holdCloudKitFetches {
 417     XCTAssertFalse([self.ckFetchHoldOperation isPending], "Shouldn't already be a pending cloudkit fetch hold operation");
 418     self.ckFetchHoldOperation = [NSBlockOperation blockOperationWithBlock:^{
 419         secnotice("ckks", "Released CloudKit fetch hold.");
 420     }];
 421 }
 422 -(void)releaseCloudKitFetchHold {
 423     if([self.ckFetchHoldOperation isPending]) {
 424         [self.operationQueue addOperation: self.ckFetchHoldOperation];
 425     }
 426 }
 427
 428 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID {
 429     [self expectCKModifyItemRecords:expectedNumberOfRecords
 430            currentKeyPointerRecords:expectedCurrentKeyRecords
 431                              zoneID:zoneID
 432                           checkItem:nil];
 433 }
 434
 435 - (void)expectCKModifyItemRecords: (NSUInteger) expectedNumberOfRecords currentKeyPointerRecords: (NSUInteger) expectedCurrentKeyRecords zoneID: (CKRecordZoneID*) zoneID checkItem: (BOOL (^)(CKRecord*)) checkItem {
 436     [self expectCKModifyItemRecords:expectedNumberOfRecords
 437                      deletedRecords:0
 438            currentKeyPointerRecords:expectedCurrentKeyRecords
 439                              zoneID:zoneID
 440                           checkItem:checkItem];
 441 }
 442
 443 - (void)expectCKModifyItemRecords:(NSUInteger)expectedNumberOfModifiedRecords
 444                    deletedRecords:(NSUInteger)expectedNumberOfDeletedRecords
 445          currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 446                            zoneID:(CKRecordZoneID*)zoneID
 447                         checkItem:(BOOL (^)(CKRecord*))checkItem {
 448     // We're updating the device state type on every update, so add it in here
 449     NSMutableDictionary* expectedRecords = [@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfModifiedRecords],
 450                                               SecCKRecordCurrentKeyType: [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords],
 451                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 452                                               } mutableCopy];
 453
 454     if(SecCKKSSyncManifests()) {
 455         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 1];
 456         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 457     }
 458
 459     NSDictionary* deletedRecords = nil;
 460     if(expectedNumberOfDeletedRecords != 0) {
 461         deletedRecords = @{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfDeletedRecords]};
 462     }
 463
 464     [self expectCKModifyRecords:expectedRecords
 465         deletedRecordTypeCounts:deletedRecords
 466                          zoneID:zoneID
 467             checkModifiedRecord: ^BOOL (CKRecord* record){
 468                 if([record.recordType isEqualToString: SecCKRecordItemType] && checkItem) {
 469                     return checkItem(record);
 470                 } else {
 471                     return YES;
 472                 }
 473             }
 474            runAfterModification:nil];
 475 }
 476
 477
 478
 479 - (void)expectCKModifyKeyRecords:(NSUInteger)expectedNumberOfRecords
 480         currentKeyPointerRecords:(NSUInteger)expectedCurrentKeyRecords
 481                  tlkShareRecords:(NSUInteger)expectedTLKShareRecords
 482                           zoneID:(CKRecordZoneID*)zoneID {
 483     NSNumber* nkeys = [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords];
 484     NSNumber* ncurrentkeys = [NSNumber numberWithUnsignedInteger: expectedCurrentKeyRecords];
 485     NSNumber* ntlkshares = [NSNumber numberWithUnsignedInteger: expectedTLKShareRecords];
 486
 487     [self expectCKModifyRecords:@{SecCKRecordIntermediateKeyType: nkeys,
 488                                   SecCKRecordCurrentKeyType: ncurrentkeys,
 489                                   SecCKRecordTLKShareType: ntlkshares,
 490                                   }
 491         deletedRecordTypeCounts:nil
 492                          zoneID:zoneID
 493             checkModifiedRecord:nil
 494            runAfterModification:nil];
 495 }
 496
 497 - (void)expectCKModifyRecords:(NSDictionary<NSString*, NSNumber*>*) expectedRecordTypeCounts
 498       deletedRecordTypeCounts:(NSDictionary<NSString*, NSNumber*>*) expectedDeletedRecordTypeCounts
 499                        zoneID:(CKRecordZoneID*) zoneID
 500           checkModifiedRecord:(BOOL (^)(CKRecord*)) checkModifiedRecord
 501          runAfterModification:(void (^) ())afterModification
 502 {
 503     __weak __typeof(self) weakSelf = self;
 504
 505     // Create an object for the block to retain and modify
 506     BoolHolder* runAlready = [[BoolHolder alloc] init];
 507
 508     secnotice("fakecloudkit", "expecting an operation matching modifications: %@ deletions: %@",
 509               expectedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 510
 511     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 512         secnotice("fakecloudkit", "Received an operation, checking");
 513         __block bool matches = false;
 514         if(runAlready.state) {
 515             secnotice("fakecloudkit", "Run already, skipping");
 516             return NO;
 517         }
 518
 519         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 520             __strong __typeof(weakSelf) strongSelf = weakSelf;
 521             XCTAssertNotNil(strongSelf, "self exists");
 522
 523             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 524             matches = true;
 525
 526             NSMutableDictionary<NSString*, NSNumber*>* modifiedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 527             NSMutableDictionary<NSString*, NSNumber*>* deletedRecordTypeCounts = [[NSMutableDictionary alloc] init];
 528
 529             // First: check if it matches. If it does, _then_ execute the operation.
 530             // Supports single-zone atomic writes only
 531
 532             if(!op.atomic) {
 533                 // We only care about atomic operations
 534                 secnotice("fakecloudkit", "Not an atomic operation; quitting: %@", op);
 535                 return NO;
 536             }
 537
 538             FakeCKZone* zone = strongSelf.zones[zoneID];
 539             XCTAssertNotNil(zone, "Have a zone for these records");
 540
 541             for(CKRecord* record in op.recordsToSave) {
 542                 if(![record.recordID.zoneID isEqual: zoneID]) {
 543                     secnotice("fakecloudkit", "Modified record zone ID mismatch: %@ %@", zoneID, record.recordID.zoneID);
 544                     return NO;
 545                 }
 546
 547                 if([zone errorFromSavingRecord: record]) {
 548                     secnotice("fakecloudkit", "Record zone rejected record write: %@", record);
 549                     return NO;
 550                 }
 551
 552                 NSNumber* currentCountNumber = modifiedRecordTypeCounts[record.recordType];
 553                 NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 554                 modifiedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 555             }
 556
 557             for(CKRecordID* recordID in op.recordIDsToDelete) {
 558                 if(![recordID.zoneID isEqual: zoneID]) {
 559                     matches = false;
 560                     secnotice("fakecloudkit", "Deleted record zone ID mismatch: %@ %@", zoneID, recordID.zoneID);
 561                 }
 562
 563                 // Find the object in CloudKit, and record its type
 564                 CKRecord* record = strongSelf.zones[zoneID].currentDatabase[recordID];
 565                 if(record) {
 566                     NSNumber* currentCountNumber = deletedRecordTypeCounts[record.recordType];
 567                     NSUInteger currentCount = currentCountNumber ? [currentCountNumber unsignedIntegerValue] : 0;
 568                     deletedRecordTypeCounts[record.recordType] = [NSNumber numberWithUnsignedInteger: currentCount + 1];
 569                 }
 570             }
 571
 572             NSMutableDictionary* filteredExpectedRecordTypeCounts = [expectedRecordTypeCounts mutableCopy];
 573             for(NSString* key in filteredExpectedRecordTypeCounts.allKeys) {
 574                 if([filteredExpectedRecordTypeCounts[key] isEqual: [NSNumber numberWithInt:0]]) {
 575                     filteredExpectedRecordTypeCounts[key] = nil;
 576                 }
 577             }
 578             filteredExpectedRecordTypeCounts[SecCKRecordManifestType] = modifiedRecordTypeCounts[SecCKRecordManifestType];
 579             filteredExpectedRecordTypeCounts[SecCKRecordManifestLeafType] = modifiedRecordTypeCounts[SecCKRecordManifestLeafType];
 580
 581             // Inspect that we have exactly the same records as we expect
 582             if(expectedRecordTypeCounts) {
 583                 matches &= !![modifiedRecordTypeCounts isEqual: filteredExpectedRecordTypeCounts];
 584                 if(!matches) {
 585                     secnotice("fakecloudkit", "Record number mismatch: %@ %@", modifiedRecordTypeCounts, filteredExpectedRecordTypeCounts);
 586                     return NO;
 587                 }
 588             } else {
 589                 matches &= op.recordsToSave.count == 0u;
 590                 if(!matches) {
 591                     secnotice("fakecloudkit", "Record number mismatch: %@ 0", modifiedRecordTypeCounts);
 592                     return NO;
 593                 }
 594             }
 595             if(expectedDeletedRecordTypeCounts) {
 596                 matches &= !![deletedRecordTypeCounts  isEqual: expectedDeletedRecordTypeCounts];
 597                 if(!matches) {
 598                     secnotice("fakecloudkit", "Deleted record number mismatch: %@ %@", deletedRecordTypeCounts, expectedDeletedRecordTypeCounts);
 599                     return NO;
 600                 }
 601             } else {
 602                 matches &= op.recordIDsToDelete.count == 0u;
 603                 if(!matches) {
 604                     secnotice("fakecloudkit", "Deleted record number mismatch: %@ 0", deletedRecordTypeCounts);
 605                     return NO;
 606                 }
 607             }
 608
 609             // We have the right number of things, and their etags match. Ensure that they have the right etags
 610             if(matches && checkModifiedRecord) {
 611                 // Clearly we have the right number of things. Call checkRecord on them...
 612                 for(CKRecord* record in op.recordsToSave) {
 613                     matches &= !!(checkModifiedRecord(record));
 614                     if(!matches) {
 615                         secnotice("fakecloudkit", "Check record reports NO: %@ 0", record);
 616                         return NO;
 617                     }
 618                 }
 619             }
 620
 621             if(matches) {
 622                 // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 623                 // if you'd like to read the data from this write.
 624                 NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-write" withBlock: ^{
 625                     @synchronized(zone.currentDatabase) {
 626                         NSMutableArray* savedRecords = [[NSMutableArray alloc] init];
 627                         for(CKRecord* record in op.recordsToSave) {
 628                             CKRecord* reflectedRecord = [record copy];
 629                             reflectedRecord.modificationDate = [NSDate date];
 630
 631                             [zone addToZone: reflectedRecord];
 632
 633                             [savedRecords addObject:reflectedRecord];
 634                             op.perRecordCompletionBlock(reflectedRecord, nil);
 635                         }
 636                         for(CKRecordID* recordID in op.recordIDsToDelete) {
 637                             // I don't believe CloudKit fails an operation if you delete a record that's not there, so:
 638                             [zone deleteCKRecordIDFromZone: recordID];
 639                         }
 640
 641                         op.modifyRecordsCompletionBlock(savedRecords, op.recordIDsToDelete, nil);
 642
 643                         if(afterModification) {
 644                             afterModification();
 645                         }
 646
 647                         op.isFinished = YES;
 648                     }
 649                 }];
 650                 [ckop addNullableDependency:strongSelf.ckModifyHoldOperation];
 651                 [strongSelf.operationQueue addOperation: ckop];
 652             }
 653         }
 654         if(matches) {
 655             runAlready.state = true;
 656         }
 657         return matches ? YES : NO;
 658     }]];
 659 }
 660
 661 - (void)failNextZoneCreation:(CKRecordZoneID*)zoneID {
 662     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 663     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 664     self.zones[zoneID].creationError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 665 }
 666
 667 // Report success, but don't actually create the zone.
 668 // This way, you can find ZoneNotFound errors later on
 669 - (void)failNextZoneCreationSilently:(CKRecordZoneID*)zoneID {
 670     XCTAssertNil(self.zones[zoneID], "Zone does not exist yet");
 671     self.zones[zoneID] = [[FakeCKZone alloc] initZone: zoneID];
 672     self.zones[zoneID].failCreationSilently = true;
 673 }
 674
 675 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID {
 676     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 677     self.zones[zoneID].subscriptionError = [[CKPrettyError alloc] initWithDomain:CKErrorDomain code:CKErrorNetworkUnavailable userInfo:@{}];
 678 }
 679
 680 - (void)failNextZoneSubscription:(CKRecordZoneID*)zoneID withError:(NSError*)error {
 681     XCTAssertNotNil(self.zones[zoneID], "Zone exists");
 682     self.zones[zoneID].subscriptionError = error;
 683 }
 684
 685 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 686     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:nil];
 687 }
 688
 689 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)())blockAfterReject {
 690     [self failNextCKAtomicModifyItemRecordsUpdateFailure:zoneID blockAfterReject:blockAfterReject withError:nil];
 691 }
 692
 693 - (void)failNextCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID blockAfterReject: (void (^)())blockAfterReject withError:(NSError*)error {
 694     __weak __typeof(self) weakSelf = self;
 695
 696     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 697         __strong __typeof(weakSelf) strongSelf = weakSelf;
 698         XCTAssertNotNil(strongSelf, "self exists");
 699
 700         __block bool rejected = false;
 701         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 702             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 703
 704             if(!op.atomic) {
 705                 // We only care about atomic operations
 706                 return NO;
 707             }
 708
 709             // We want to only match zone updates pertaining to this zone
 710             for(CKRecord* record in op.recordsToSave) {
 711                 if(![record.recordID.zoneID isEqual: zoneID]) {
 712                     return NO;
 713                 }
 714             }
 715
 716             FakeCKZone* zone = strongSelf.zones[zoneID];
 717             XCTAssertNotNil(zone, "Have a zone for these records");
 718
 719             rejected = true;
 720
 721             if(error) {
 722                 [strongSelf rejectWrite: op withError:error];
 723             } else {
 724                 NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 725                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 726             }
 727
 728             if(blockAfterReject) {
 729                 blockAfterReject();
 730             }
 731         }
 732         return rejected ? YES : NO;
 733     }]];
 734 }
 735
 736 - (void)expectCKAtomicModifyItemRecordsUpdateFailure:(CKRecordZoneID*)zoneID {
 737     __weak __typeof(self) weakSelf = self;
 738
 739     [[self.mockDatabase expect] addOperation:[OCMArg checkWithBlock:^BOOL(id obj) {
 740         __strong __typeof(weakSelf) strongSelf = weakSelf;
 741         XCTAssertNotNil(strongSelf, "self exists");
 742
 743         __block bool rejected = false;
 744         if ([obj isKindOfClass:[CKModifyRecordsOperation class]]) {
 745             CKModifyRecordsOperation *op = (CKModifyRecordsOperation *)obj;
 746
 747             if(!op.atomic) {
 748                 // We only care about atomic operations
 749                 return NO;
 750             }
 751
 752             // We want to only match zone updates pertaining to this zone
 753             for(CKRecord* record in op.recordsToSave) {
 754                 if(![record.recordID.zoneID isEqual: zoneID]) {
 755                     return NO;
 756                 }
 757             }
 758
 759             FakeCKZone* zone = strongSelf.zones[zoneID];
 760             XCTAssertNotNil(zone, "Have a zone for these records");
 761
 762             NSMutableDictionary<CKRecordID*, NSError*>* failedRecords = [[NSMutableDictionary alloc] init];
 763
 764             @synchronized(zone.currentDatabase) {
 765                 for(CKRecord* record in op.recordsToSave) {
 766                     // Check if we should allow this transaction
 767                     NSError* recordSaveError = [zone errorFromSavingRecord: record];
 768                     if(recordSaveError) {
 769                         failedRecords[record.recordID] = recordSaveError;
 770                         rejected = true;
 771                     }
 772                 }
 773             }
 774
 775             if(rejected) {
 776                 [strongSelf rejectWrite: op failedRecords:failedRecords];
 777             }
 778         }
 779         return rejected ? YES : NO;
 780     }]];
 781 }
 782
 783 -(void)rejectWrite:(CKModifyRecordsOperation*)op withError:(NSError*)error {
 784     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 785     // if you'd like to read the data from this write.
 786     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write-error" withBlock: ^{
 787         op.modifyRecordsCompletionBlock(nil, nil, error);
 788         op.isFinished = YES;
 789     }];
 790     [ckop addNullableDependency: self.ckModifyHoldOperation];
 791     [self.operationQueue addOperation: ckop];
 792 }
 793
 794 -(void)rejectWrite:(CKModifyRecordsOperation*)op failedRecords:(NSMutableDictionary<CKRecordID*, NSError*>*)failedRecords {
 795     // Add the batch request failed errors
 796     for(CKRecord* record in op.recordsToSave) {
 797         NSError* exists = failedRecords[record.recordID];
 798         if(!exists) {
 799             // TODO: might have important userInfo, but we're not mocking that yet
 800             failedRecords[record.recordID] = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorBatchRequestFailed userInfo: @{}];
 801         }
 802     }
 803
 804     NSError* error = [[CKPrettyError alloc] initWithDomain: CKErrorDomain code: CKErrorPartialFailure userInfo: @{CKPartialErrorsByItemIDKey: failedRecords}];
 805
 806     // Emulate cloudkit and schedule the operation for execution. Be sure to wait for this operation
 807     // if you'd like to read the data from this write.
 808     NSBlockOperation* ckop = [NSBlockOperation named:@"cloudkit-reject-write" withBlock: ^{
 809         op.modifyRecordsCompletionBlock(nil, nil, error);
 810         op.isFinished = YES;
 811     }];
 812     [ckop addNullableDependency: self.ckModifyHoldOperation];
 813     [self.operationQueue addOperation: ckop];
 814 }
 815
 816 - (void)expectCKDeleteItemRecords:(NSUInteger)expectedNumberOfRecords
 817                            zoneID:(CKRecordZoneID*) zoneID {
 818
 819     // We're updating the device state type on every update, so add it in here
 820     NSMutableDictionary* expectedRecords = [@{
 821                                               SecCKRecordDeviceStateType: [NSNumber numberWithUnsignedInt: 1],
 822                                               } mutableCopy];
 823     if(SecCKKSSyncManifests()) {
 824         // TODO: this really shouldn't be 2.
 825         expectedRecords[SecCKRecordManifestType] = [NSNumber numberWithInt: 2];
 826         expectedRecords[SecCKRecordManifestLeafType] = [NSNumber numberWithInt: 72];
 827     }
 828
 829     [self expectCKModifyRecords:expectedRecords
 830         deletedRecordTypeCounts:@{SecCKRecordItemType: [NSNumber numberWithUnsignedInteger: expectedNumberOfRecords]}
 831                          zoneID:zoneID
 832             checkModifiedRecord:nil
 833            runAfterModification:nil];
 834 }
 835
 836 -(void)waitForCKModifications {
 837     // CloudKit modifications are put on the local queue.
 838     // This is heavyweight but should suffice.
 839     [self.operationQueue waitUntilAllOperationsAreFinished];
 840 }
 841
 842 - (void)tearDown {
 843     // Put teardown code here. This method is called after the invocation of each test method in the class.
 844
 845     if(SecCKKSIsEnabled()) {
 846         // Ensure we don't have any blocking operations
 847         self.accountStatus = CKAccountStatusNoAccount;
 848         [self startCKKSSubsystem];
 849
 850         [self waitForCKModifications];
 851
 852         XCTAssertEqual(0, [self.injectedManager.completedSecCKKSInitialize wait:2*NSEC_PER_SEC],
 853             "Timeout did not occur waiting for SecCKKSInitialize");
 854
 855         // Make sure this happens before teardown.
 856         XCTAssertEqual(0, [self.accountStateTracker.finishedInitialCalls wait:1*NSEC_PER_SEC], "Account state tracker initialized itself");
 857     }
 858
 859     [super tearDown];
 860
 861     [self.injectedManager cancelPendingOperations];
 862     [CKKSViewManager resetManager:true setTo:nil];
 863     self.injectedManager = nil;
 864     [self.mockCKKSViewManager stopMocking];
 865     self.mockCKKSViewManager = nil;
 866
 867     [self.mockAccountStateTracker stopMocking];
 868     self.mockAccountStateTracker = nil;
 869
 870     [self.mockLockStateTracker stopMocking];
 871     self.mockLockStateTracker = nil;
 872
 873     [self.mockFakeCKModifyRecordZonesOperation stopMocking];
 874     self.mockFakeCKModifyRecordZonesOperation = nil;
 875
 876     [self.mockFakeCKModifySubscriptionsOperation stopMocking];
 877     self.mockFakeCKModifySubscriptionsOperation = nil;
 878
 879     [self.mockFakeCKFetchRecordZoneChangesOperation stopMocking];
 880     self.mockFakeCKFetchRecordZoneChangesOperation = nil;
 881
 882     [self.mockFakeCKFetchRecordsOperation stopMocking];
 883     self.mockFakeCKFetchRecordsOperation = nil;
 884
 885     [self.mockFakeCKQueryOperation stopMocking];
 886     self.mockFakeCKQueryOperation = nil;
 887
 888     [self.mockDatabase stopMocking];
 889     self.mockDatabase = nil;
 890
 891     [self.mockContainer stopMocking];
 892     self.mockContainer = nil;
 893
 894     self.zones = nil;
 895     self.operationQueue = nil;
 896     self.ckksHoldOperation = nil;
 897     self.ckaccountHoldOperation = nil;
 898
 899     SecCKKSTestResetFlags();
 900 }
 901
 902 - (CKKSKey*) fakeTLK: (CKRecordZoneID*)zoneID {
 903     CKKSKey* key = [[CKKSKey alloc] initSelfWrappedWithAESKey:[[CKKSAESSIVKey alloc] initWithBase64: @"uImdbZ7Zg+6WJXScTnRBfNmoU1UiMkSYxWc+d1Vuq3IFn2RmTRkTdWTe3HmeWo1pAomqy+upK8KHg2PGiRGhqg=="]
 904                                                          uuid:[[NSUUID UUID] UUIDString]
 905                                                      keyclass:SecCKKSKeyClassTLK
 906                                                         state: SecCKKSProcessedStateLocal
 907                                                        zoneID:zoneID
 908                                               encodedCKRecord: nil
 909                                                    currentkey: true];
 910     [key CKRecordWithZoneID: zoneID];
 911     return key;
 912 }
 913
 914 - (NSError*)ckInternalServerExtensionError:(NSInteger)code description:(NSString*)desc {
 915     NSError* extensionError = [[CKPrettyError alloc] initWithDomain:@"CloudkitKeychainService"
 916                                                                code:code
 917                                                            userInfo:@{
 918                                                                       CKErrorServerDescriptionKey: desc,
 919                                                                       NSLocalizedDescriptionKey: desc,
 920                                                                       }];
 921     NSError* internalError = [[CKPrettyError alloc] initWithDomain:CKInternalErrorDomain
 922                                                               code:CKErrorInternalPluginError
 923                                                           userInfo:@{CKErrorServerDescriptionKey: desc,
 924                                                                      NSLocalizedDescriptionKey: desc,
 925                                                                      NSUnderlyingErrorKey: extensionError,
 926                                                                      }];
 927     NSError* error = [[CKPrettyError alloc] initWithDomain:CKErrorDomain
 928                                                       code:CKErrorServerRejectedRequest
 929                                                   userInfo:@{NSUnderlyingErrorKey: internalError,
 930                                                              CKErrorServerDescriptionKey: desc,
 931                                                              NSLocalizedDescriptionKey: desc,
 932                                                              CKContainerIDKey: SecCKKSContainerName,
 933                                                              }];
 934     return error;
 935 }
 936
 937 @end
 938
 939 #endif