Merge f66b64c into c110245

google · Sep 24, 2021 · ae31688 · ae31688
2 parents c110245 + f66b64c
commit ae31688
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 62 deletions.
diff --git 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
@@ -100,7 +100,10 @@ - (instancetype)init {
   self = [super init];
   if (self) {
     _responseCallbacks = [NSMutableArray array];
-    _subscribed = YES;
+    _subscriptions = [NSMutableArray arrayWithCapacity:ES_EVENT_TYPE_LAST];
+    for (size_t i = 0; i < ES_EVENT_TYPE_LAST; i++) {
+      _subscriptions[i] = @NO;
+    }
   }
   return self;
 };
@@ -110,7 +113,6 @@ - (void)reset {
     [self.responseCallbacks removeAllObjects];
     self.handler = nil;
     self.client = nil;
-    self.subscribed = NO;
   }
 };
 
@@ -194,14 +196,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;
+  for (size_t i = 0; i < event_count; i++) {
+    [MockEndpointSecurity mockEndpointSecurity].subscriptions[events[i]] = @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;
+  for (size_t i = 0; i < event_count; i++) {
+    [MockEndpointSecurity mockEndpointSecurity].subscriptions[events[i]] = @NO;
+  }
 
   return ES_RETURN_SUCCESS;
 };
diff --git 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
@@ -50,36 +50,25 @@ - (void)checkBinaryExecution:(NSString *)binaryName
   OCMStub([self.mockSNTDatabaseController databasePath]).andReturn(testPath);
 
   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])
       ;
     [santaInit fulfill];
   });
 
-  [self waitForExpectationsWithTimeout:30.0
-                               handler:^(NSError *error) {
-                                 if (error) {
-                                   XCTFail(@"Santa's subscription to EndpointSecurity timed out "
-                                           @"with error: %@",
-                                           error);
-                                 }
-                               }];
+  [app start];
+  [self waitForExpectations:@[ santaInit ] timeout:60.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 +84,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:60.0];
 
   XCTAssertEqual(got.result, wantResult, @"received unexpected ES response on executing \"%@/%@\"",
                  testPath, binaryName);