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