Change mock es_subscribe to note specific events we're ready for

google · Sep 24, 2021 · f66b64c · f66b64c
1 parent dcf726b
commit f66b64c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 9 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/SNTApplicationTest.m b/Source/santad/SNTApplicationTest.m
@@ -50,19 +50,17 @@ - (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];
   });
 
+  [app start];
   [self waitForExpectations:@[ santaInit ] timeout:60.0];
 
   XCTestExpectation *expectation =