santa-driver: Fix race condition by adding CAS op to SantaCache

Change the signature of the set method in SantaCache so that it takes an
optional previous-value parameter (and a bool indicating that this value
has been provided). If previous-value is provided, set becomes a
compare-and-swap. Also provide 2 overloads for a cleaner interface, one
with and without the previous-value parameter.

Source/santa-driver/SantaCache.h

@@ -101,12 +101,19 @@ template<class T> class SantaCache {
   /**
     Set an element in the cache.
 
-    @note If the cache is full when this is called, this will empty the cache before
-    inserting the new value.
+    @note If the cache is full when this is called, this will
+          empty the cache before inserting the new value.
 
-    @return if an existing value was replaced, the previous value, otherwise zero_
+    @param key, The key
+    @param value, The value with parameterized type
+    @param previous_value, If the has_prev_value parameter is true the new
+        value will only be set if this parameter is equal to the provided value.
+        This allows set to become a CAS operation.
+    @param has_prev_value, Pass true if previous_value should be used.
+
+    @return the previous value (which may be zero_)
   */
-  T set(uint64_t key, T value) {
+  T set(uint64_t key, T value, T previous_value, bool has_prev_value) {
     struct bucket *bucket = &buckets_[hash(key)];
     lock(bucket);
     struct entry *entry = (struct entry *)((uintptr_t)bucket->head - 1);
@@ -114,6 +121,12 @@ template<class T> class SantaCache {
     while (entry != nullptr) {
       if (entry->key == key) {
         T existing_value = entry->value;
+
+        if (has_prev_value && previous_value != existing_value) {
+          unlock(bucket);
+          return existing_value;
+        }
+
         entry->value = value;
 
         if (value == zero_) {
@@ -134,13 +147,15 @@ template<class T> class SantaCache {
     }
 
     // If value is zero_, we're clearing but there's nothing to clear
-    // so we don't need to do anything else.
-    if (value == zero_) {
+    // so we don't need to do anything else. Alternatively, if has_prev_value
+    // is true and is not zero_ we don't want to set a value.
+    if (value == zero_ || (has_prev_value && previous_value != zero_)) {
       unlock(bucket);
       return zero_;
     }
 
-    // Check that adding this new item won't take the cache over its maximum size.
+    // Check that adding this new item won't take the cache
+    // over its maximum size.
     if (count_ + 1 > max_size_) {
       unlock(bucket);
       lock(&clear_bucket_);
@@ -152,9 +167,10 @@ template<class T> class SantaCache {
       unlock(&clear_bucket_);
     }
 
-    // Allocate a new entry, set the key and value, then set the next pointer as the current
-    // first entry in the bucket then make this new entry the first in the bucket.
-    struct entry *new_entry = (struct entry *)IOMallocAligned(sizeof(struct entry), 2);
+    // Allocate a new entry, set the key and value, then put this new entry at
+    // the head of this bucket's linked list.
+    struct entry *new_entry = (struct entry *)IOMallocAligned(
+        sizeof(struct entry), 2);
     new_entry->key = key;
     new_entry->value = value;
     new_entry->next = (struct entry *)((uintptr_t)bucket->head - 1);
@@ -165,6 +181,20 @@ template<class T> class SantaCache {
     return zero_;
   }
 
+  /**
+    Overload to allow setting without providing a previous value
+  */
+  T set(uint64_t key, T value) {
+    return set(key, value, {}, false);
+  }
+
+  /**
+    Overload to allow setting while providing a previous value
+  */
+  T set(uint64_t key, T value, T previous_value) {
+    return set(key, value, previous_value, true);
+  }
+
   /**
     An alias for `set(key, zero_)`
   */
@@ -247,7 +277,7 @@ template<class T> class SantaCache {
   /**
     Holder for a 'zero' entry for the current type
   */
-  const T zero_ = T();
+  const T zero_ = {};
 
   /**
     Special bucket used when automatically clearing due to size

Source/santa-driver/SantaDecisionManager.cc

@@ -179,7 +179,8 @@ kern_return_t SantaDecisionManager::StartListener() {
   if (!vnode_listener_) return kIOReturnInternalError;
 
   fileop_listener_ = kauth_listen_scope(
-      KAUTH_SCOPE_FILEOP, fileop_scope_callback, reinterpret_cast<void *>(this));
+      KAUTH_SCOPE_FILEOP, fileop_scope_callback,
+      reinterpret_cast<void *>(this));
   if (!fileop_listener_) return kIOReturnInternalError;
 
   LOGD("Listeners started.");
@@ -215,8 +216,10 @@ kern_return_t SantaDecisionManager::StopListener() {
   @param identifier The identifier
   @return SantaCache* The cache to use
 */
-SantaCache<uint64_t>* SantaDecisionManager::CacheForIdentifier(const uint64_t identifier) {
-  return (identifier >> 32 == root_fsid_) ? root_decision_cache_ : non_root_decision_cache_;
+SantaCache<uint64_t>* SantaDecisionManager::CacheForIdentifier(
+    const uint64_t identifier) {
+  return (identifier >> 32 == root_fsid_) ?
+    root_decision_cache_ : non_root_decision_cache_;
 }
 
 void SantaDecisionManager::AddToCache(
@@ -226,10 +229,17 @@ void SantaDecisionManager::AddToCache(
 
   auto decision_cache = CacheForIdentifier(identifier);
 
-  // If a previous entry was not found and the new entry is not `REQUEST_BINARY`, remove the
-  // existing entry. This is to prevent adding an ALLOW to the cache after a write has occurred.
-  if (decision_cache->set(identifier, val) == 0 && decision != ACTION_REQUEST_BINARY) {
-    decision_cache->remove(identifier);
+  switch (decision) {
+    case ACTION_REQUEST_BINARY:
+      decision_cache->set(identifier, val, 0);
+      break;
+    case ACTION_RESPOND_ALLOW:
+    case ACTION_RESPOND_DENY:
+      decision_cache->set(
+          identifier, val, ((uint64_t)ACTION_REQUEST_BINARY << 56));
+      break;
+    default:
+      break;
   }
 
   wakeup((void *)identifier);
@@ -282,7 +292,8 @@ santa_action_t SantaDecisionManager::GetFromCache(uint64_t identifier) {
   return result;
 }
 
-santa_action_t SantaDecisionManager::GetFromDaemon(santa_message_t *message, uint64_t identifier) {
+santa_action_t SantaDecisionManager::GetFromDaemon(
+    santa_message_t *message, uint64_t identifier) {
   auto return_action = ACTION_UNSET;
 
 #ifdef DEBUG
@@ -294,10 +305,11 @@ santa_action_t SantaDecisionManager::GetFromDaemon(santa_message_t *message, uin
 
   // Wait for the daemon to respond or die.
   do {
-    // Add pending request to cache, to be replaced by daemon with actual response
+    // Add pending request to cache, to be replaced
+    // by daemon with actual response.
     AddToCache(identifier, ACTION_REQUEST_BINARY, 0);
 
-    // Send request to daemon...
+    // Send request to daemon.
     if (!PostToDecisionQueue(message)) {
       LOGE("Failed to queue request for %s.", message->path);
       RemoveFromCache(identifier);
@@ -343,6 +355,8 @@ santa_action_t SantaDecisionManager::FetchDecision(
     if (RESPONSE_VALID(return_action)) {
       return return_action;
     } else if (return_action == ACTION_REQUEST_BINARY) {
+      // This thread will now sleep for kRequestLoopSleepMilliseconds (1s) or
+      // until AddToCache is called, indicating a response has arrived.
       msleep((void *)vnode_id, NULL, 0, "", &ts_);
     } else {
       break;

Tests/LogicTests/SantaCacheTest.mm

@@ -150,4 +150,27 @@
   delete sut;
 }
 
+- (void)testCompareAndSwap {
+  auto sut = new SantaCache<uint64_t>(100, 2);
+
+  sut->set(1, 42);
+  sut->set(1, 666, 1);
+  sut->set(1, 666, 0);
+  XCTAssertEqual(sut->get(1), 42);
+
+  sut->set(1, 0);
+  XCTAssertEqual(sut->get(1), 0);
+
+  sut->set(1, 42, 1);
+  XCTAssertEqual(sut->get(1), 0);
+
+  sut->set(1, 42, 0);
+  XCTAssertEqual(sut->get(1), 42);
+
+  sut->set(1, 0, 666);
+  XCTAssertEqual(sut->get(1), 42);
+  sut->set(1, 0, 42);
+  XCTAssertEqual(sut->get(1), 0);
+}
+
 @end