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