diff --git a/Source/santad/EventProviders/EndpointSecurityTestUtil.h b/Source/santad/EventProviders/EndpointSecurityTestUtil.h index b04c06fd0..7986e36ee 100644 --- a/Source/santad/EventProviders/EndpointSecurityTestUtil.h +++ b/Source/santad/EventProviders/EndpointSecurityTestUtil.h @@ -44,7 +44,7 @@ typedef void (^ESCallback)(ESResponse *_Nonnull); // Singleton wrapper around all of the kernel-level EndpointSecurity framework functions. @interface MockEndpointSecurity : NSObject -@property BOOL subscribed; +@property NSMutableArray *_Nonnull subscriptions; - (void)reset; - (void)registerResponseCallback:(ESCallback _Nonnull)callback; - (void)triggerHandler:(es_message_t *_Nonnull)msg; diff --git a/Source/santad/EventProviders/EndpointSecurityTestUtil.mm b/Source/santad/EventProviders/EndpointSecurityTestUtil.mm index bbe882404..66d9258a5 100644 --- a/Source/santad/EventProviders/EndpointSecurityTestUtil.mm +++ b/Source/santad/EventProviders/EndpointSecurityTestUtil.mm @@ -100,17 +100,23 @@ - (instancetype)init { self = [super init]; if (self) { _responseCallbacks = [NSMutableArray array]; - _subscribed = YES; + _subscriptions = [NSMutableArray arrayWithCapacity:ES_EVENT_TYPE_LAST]; + [self resetSubscriptions]; } return self; }; +- (void)resetSubscriptions { + for (size_t i = 0; i < ES_EVENT_TYPE_LAST; i++) { + _subscriptions[i] = @NO; + } +} + - (void)reset { @synchronized(self) { [self.responseCallbacks removeAllObjects]; self.handler = nil; self.client = nil; - self.subscribed = NO; } }; @@ -148,6 +154,16 @@ - (es_respond_result_t)respond_auth_result:(const es_message_t *_Nonnull)msg return ES_RESPOND_RESULT_SUCCESS; }; +- (void)setSubscriptions:(const es_event_type_t *_Nonnull)events + event_count:(uint32_t)event_count + value:(NSNumber *)value { + @synchronized(self) { + for (size_t i = 0; i < event_count; i++) { + self.subscriptions[events[i]] = value; + } + } +} + + (instancetype _Nonnull)mockEndpointSecurity { static MockEndpointSecurity *sharedES; static dispatch_once_t onceToken; @@ -194,14 +210,18 @@ es_respond_result_t es_respond_auth_result(es_client_t *_Nonnull client, API_UNAVAILABLE(ios, tvos, watchos) es_return_t es_subscribe(es_client_t *_Nonnull client, const es_event_type_t *_Nonnull events, uint32_t event_count) { - [MockEndpointSecurity mockEndpointSecurity].subscribed = YES; + [[MockEndpointSecurity mockEndpointSecurity] setSubscriptions:events + event_count:event_count + value:@YES]; return ES_RETURN_SUCCESS; } API_AVAILABLE(macos(10.15)) API_UNAVAILABLE(ios, tvos, watchos) es_return_t es_unsubscribe(es_client_t *_Nonnull client, const es_event_type_t *_Nonnull events, uint32_t event_count) { - [MockEndpointSecurity mockEndpointSecurity].subscribed = NO; + [[MockEndpointSecurity mockEndpointSecurity] setSubscriptions:events + event_count:event_count + value:@NO]; return ES_RETURN_SUCCESS; }; diff --git a/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm b/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm index ef6fe3955..9a293bb18 100644 --- a/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm +++ b/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm @@ -71,14 +71,7 @@ - (void)testDenyOnTimeout { [mockES triggerHandler:m.message]; - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa auth test timed out without receiving two " - @"events. Instead, had error: %@", - error); - } - }]; + [self waitForExpectations:@[ expectation ] timeout:60.0]; for (ESResponse *resp in events) { XCTAssertEqual( @@ -116,12 +109,7 @@ - (void)testDeleteRulesDB { [mockES triggerHandler:m.message]; - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa auth test timed out with error: %@", error); - } - }]; + [self waitForExpectations:@[ expectation ] timeout:60.0]; XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue, @"Incorrect handling of delete of %@", testPath); @@ -152,12 +140,8 @@ - (void)testSkipOtherESEvents { }]; [mockES triggerHandler:m.message]; - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa auth test timed out with error: %@", error); - } - }]; + + [self waitForExpectations:@[ expectation ] timeout:60.0]; XCTAssertEqual(got.result, ES_AUTH_RESULT_ALLOW); } @@ -199,12 +183,7 @@ - (void)testRenameOverwriteRulesDB { [mockES triggerHandler:m.message]; - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa auth test timed out with error: %@", error); - } - }]; + [self waitForExpectations:@[ expectation ] timeout:60.0]; XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue, @"Incorrect handling of rename of %@", testPath); @@ -254,12 +233,8 @@ - (void)testRenameRulesDB { [mockES triggerHandler:m.message]; - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa auth test timed out with error: %@", error); - } - }]; + [self waitForExpectations:@[ expectation ] timeout:60.0]; + XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue, @"Incorrect handling of rename of %@", testPath); diff --git a/Source/santad/SNTApplicationTest.m b/Source/santad/SNTApplicationTest.m index addf4ccd8..49b8be8b5 100644 --- a/Source/santad/SNTApplicationTest.m +++ b/Source/santad/SNTApplicationTest.m @@ -52,34 +52,26 @@ - (void)checkBinaryExecution:(NSString *)binaryName SNTApplication *app = [[SNTApplication alloc] init]; [app start]; - // es events will start flowing in as soon as es_subscribe is called, regardless - // of whether we're ready or not for it. XCTestExpectation *santaInit = [self expectationWithDescription:@"Wait for Santa to subscribe to EndpointSecurity"]; dispatch_async(dispatch_get_global_queue(QOS_CLASS_BACKGROUND, 0), ^{ - while (!mockES.subscribed) + while ([mockES.subscriptions[ES_EVENT_TYPE_AUTH_EXEC] isEqualTo:@NO]) ; [santaInit fulfill]; }); - [self waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail(@"Santa's subscription to EndpointSecurity timed out " - @"with error: %@", - error); - } - }]; + // Ugly hack to deflake the test and allow listenForDecisionRequests to install the correct + // decision callback. + sleep(1); + [self waitForExpectations:@[ santaInit ] timeout:10.0]; XCTestExpectation *expectation = [self expectationWithDescription:@"Wait for santa's Auth dispatch queue"]; __block ESResponse *got = nil; [mockES registerResponseCallback:^(ESResponse *r) { - @synchronized(self) { - got = r; - [expectation fulfill]; - } + got = r; + [expectation fulfill]; }]; NSString *binaryPath = [NSString pathWithComponents:@[ testPath, binaryName ]]; @@ -95,15 +87,7 @@ - (void)checkBinaryExecution:(NSString *)binaryName [mockES triggerHandler:msg.message]; - [self - waitForExpectationsWithTimeout:30.0 - handler:^(NSError *error) { - if (error) { - XCTFail( - @"Santa auth test on binary \"%@/%@\" timed out with error: %@", - testPath, binaryName, error); - } - }]; + [self waitForExpectations:@[ expectation ] timeout:10.0]; XCTAssertEqual(got.result, wantResult, @"received unexpected ES response on executing \"%@/%@\"", testPath, binaryName);