Fix automatically denied events with small deadlines

Source/santad/BUILD

@@ -1182,6 +1182,7 @@ santa_unit_test(
         ":SNTEndpointSecurityClient",
         ":WatchItemPolicy",
         "//Source/common:SNTConfigurator",
+        "//Source/common:SystemResources",
         "//Source/common:TestUtils",
         "@OCMock",
         "@com_google_googletest//:gtest",

Source/santad/EventProviders/SNTEndpointSecurityClient.mm

@@ -21,6 +21,7 @@
 #include <stdlib.h>
 #include <sys/qos.h>
 
+#include <algorithm>
 #include <set>
 #include <string>
 #include <string_view>
@@ -48,7 +49,9 @@
                                                 "/private/var/db/santa/events.db"};
 
 @interface SNTEndpointSecurityClient ()
-@property int64_t deadlineMarginMS;
+@property(nonatomic) double defaultBudget;
+@property(nonatomic) int64_t minAllowedHeadroom;
+@property(nonatomic) int64_t maxAllowedHeadroom;
 @property SNTConfigurator *configurator;
 @end
 
@@ -68,10 +71,18 @@ - (instancetype)initWithESAPI:(std::shared_ptr<EndpointSecurityAPI>)esApi
   if (self) {
     _esApi = std::move(esApi);
     _metrics = std::move(metrics);
-    _deadlineMarginMS = 5000;
     _configurator = [SNTConfigurator configurator];
     _processor = processor;
 
+    // Default event processing budget is 80% of the deadline time
+    _defaultBudget = 0.8;
+
+    // For events with small deadlines, clamp processing budget to 1s headroom
+    _minAllowedHeadroom = 1 * NSEC_PER_SEC;
+
+    // For events with large deadlines, clamp processing budget to 5s headroom
+    _maxAllowedHeadroom = 5 * NSEC_PER_SEC;
+
     _authQueue = dispatch_queue_create(
       "com.google.santa.daemon.auth_queue",
       dispatch_queue_attr_make_with_qos_class(DISPATCH_QUEUE_CONCURRENT_WITH_AUTORELEASE_POOL,
@@ -255,6 +266,24 @@ - (void)asynchronouslyProcess:(Message)msg handler:(void (^)(Message &&))message
   });
 }
 
+- (int64_t)computeBudgetForDeadline:(uint64_t)deadline currentTime:(uint64_t)currentTime {
+  // First get how much time we have left
+  int64_t nanosUntilDeadline = (int64_t)MachTimeToNanos(deadline - currentTime);
+
+  // Compute the desired budget
+  int64_t budget = nanosUntilDeadline * self.defaultBudget;
+
+  // See how much headroom is left
+  int64_t headroom = nanosUntilDeadline - budget;
+
+  // Clamp headroom to maximize budget but ensure it's not so large as to not leave
+  // enough time to respond in an emergency.
+  headroom = std::clamp(headroom, self.minAllowedHeadroom, self.maxAllowedHeadroom);
+
+  // Return the processing budget given the allotted headroom
+  return nanosUntilDeadline - headroom;
+}
+
 - (void)processMessage:(Message &&)msg handler:(void (^)(const Message &))messageHandler {
   if (unlikely(msg->action_type != ES_ACTION_TYPE_AUTH)) {
     // This is a programming error
@@ -270,9 +299,8 @@ - (void)processMessage:(Message &&)msg handler:(void (^)(const Message &))messag
   dispatch_semaphore_signal(processingSema);
   dispatch_semaphore_t deadlineExpiredSema = dispatch_semaphore_create(0);
 
-  const uint64_t timeout = NSEC_PER_MSEC * (self.deadlineMarginMS);
-
-  uint64_t deadlineNano = MachTimeToNanos(msg->deadline - mach_absolute_time());
+  int64_t processingBudget = [self computeBudgetForDeadline:msg->deadline
+                                                currentTime:mach_absolute_time()];
 
   // TODO(mlw): How should we handle `deadlineNano <= timeout`. Will currently
   // result in the deadline block being dispatched immediately (and therefore
@@ -282,21 +310,20 @@ - (void)processMessage:(Message &&)msg handler:(void (^)(const Message &))messag
   __block Message processMsg = msg;
   __block Message deadlineMsg = msg;
 
-  dispatch_after(
-    dispatch_time(DISPATCH_TIME_NOW, deadlineNano - timeout), self->_authQueue, ^(void) {
-      if (dispatch_semaphore_wait(processingSema, DISPATCH_TIME_NOW) != 0) {
-        // Handler has already responded, nothing to do.
-        return;
-      }
-
-      bool res = [self respondToMessage:deadlineMsg
-                         withAuthResult:ES_AUTH_RESULT_DENY
-                              cacheable:false];
-
-      LOGE(@"SNTEndpointSecurityClient: deadline reached: deny pid=%d, event type: %d ret=%d",
-           audit_token_to_pid(deadlineMsg->process->audit_token), deadlineMsg->event_type, res);
-      dispatch_semaphore_signal(deadlineExpiredSema);
-    });
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, processingBudget), self->_authQueue, ^(void) {
+    if (dispatch_semaphore_wait(processingSema, DISPATCH_TIME_NOW) != 0) {
+      // Handler has already responded, nothing to do.
+      return;
+    }
+
+    bool res = [self respondToMessage:deadlineMsg
+                       withAuthResult:ES_AUTH_RESULT_DENY
+                            cacheable:false];
+
+    LOGE(@"SNTEndpointSecurityClient: deadline reached: deny pid=%d, event type: %d ret=%d",
+         audit_token_to_pid(deadlineMsg->process->audit_token), deadlineMsg->event_type, res);
+    dispatch_semaphore_signal(deadlineExpiredSema);
+  });
 
   dispatch_async(self->_authQueue, ^{
     messageHandler(processMsg);

Source/santad/EventProviders/SNTEndpointSecurityClientTest.mm

@@ -23,6 +23,7 @@
 #include <memory>
 
 #import "Source/common/SNTConfigurator.h"
+#import "Source/common/SystemResources.h"
 #include "Source/common/TestUtils.h"
 #include "Source/santad/DataLayer/WatchItemPolicy.h"
 #include "Source/santad/EventProviders/EndpointSecurity/Client.h"
@@ -48,8 +49,11 @@ - (NSString *)errorMessageForNewClientResult:(es_new_client_result_t)result;
 - (void)handleMessage:(Message &&)esMsg
    recordEventMetrics:(void (^)(santa::santad::EventDisposition disposition))recordEventMetrics;
 - (BOOL)shouldHandleMessage:(const Message &)esMsg;
+- (int64_t)computeBudgetForDeadline:(uint64_t)deadline currentTime:(uint64_t)currentTime;
 
-@property int64_t deadlineMarginMS;
+@property(nonatomic) double defaultBudget;
+@property(nonatomic) int64_t minAllowedHeadroom;
+@property(nonatomic) int64_t maxAllowedHeadroom;
 @end
 
 @interface SNTEndpointSecurityClientTest : XCTestCase
@@ -503,6 +507,46 @@ - (void)testProcessMessageHandler {
   XCTBubbleMockVerifyAndClearExpectations(mockESApi.get());
 }
 
+- (void)testComputeBudgetForDeadlineCurrentTime {
+  auto mockESApi = std::make_shared<MockEndpointSecurityAPI>();
+
+  SNTEndpointSecurityClient *client =
+    [[SNTEndpointSecurityClient alloc] initWithESAPI:mockESApi
+                                             metrics:nullptr
+                                           processor:Processor::kUnknown];
+
+  // The test uses crafted values to make even numbers. Ensure the client has
+  // expected values for these properties so the test can fail early if not.
+  XCTAssertEqual(client.defaultBudget, 0.8);
+  XCTAssertEqual(client.minAllowedHeadroom, 1 * NSEC_PER_SEC);
+  XCTAssertEqual(client.maxAllowedHeadroom, 5 * NSEC_PER_SEC);
+
+  std::map<uint64_t, int64_t> deadlineMillisToBudgetMillis{
+    // Further out deadlines clamp processing budget to maxAllowedHeadroom
+    {45000, 40000},
+
+    // Closer deadlines allow a set percentage processing budget
+    {15000, 12000},
+
+    // Near deadlines clamp processing budget to minAllowedHeadroom
+    {3500, 2500}};
+
+  uint64_t curTime = mach_absolute_time();
+
+  for (const auto [deadlineMS, budgetMS] : deadlineMillisToBudgetMillis) {
+    int64_t got =
+      [client computeBudgetForDeadline:AddNanosecondsToMachTime(deadlineMS * NSEC_PER_MSEC, curTime)
+                           currentTime:curTime];
+
+    // Add 100us, then clip to ms to account for non-exact values due to timebase division
+    got = (int64_t)((double)(got + (100 * NSEC_PER_USEC)) / (double)NSEC_PER_MSEC);
+
+    XCTAssertEqual(got, budgetMS);
+  }
+
+  XCTBubbleMockVerifyAndClearExpectations(mockESApi.get());
+}
+
 - (void)testProcessMessageHandlerWithDeadlineTimeout {
   // Set a es_message_t deadline of 750ms
   // Set a deadline leeway in the `SNTEndpointSecurityClient` of 500ms
@@ -537,7 +581,10 @@ - (void)testProcessMessageHandlerWithDeadlineTimeout {
     [[SNTEndpointSecurityClient alloc] initWithESAPI:mockESApi
                                              metrics:nullptr
                                            processor:Processor::kUnknown];
-  client.deadlineMarginMS = 500;
+
+  // Set min/max headroom the same to clamp the value for this test
+  client.minAllowedHeadroom = 500 * NSEC_PER_MSEC;
+  client.maxAllowedHeadroom = 500 * NSEC_PER_MSEC;
 
   {
     __block long result;