diff --git a/Source/santa-driver/SantaCache.h b/Source/santa-driver/SantaCache.h index 6fcfbc85..ef78e708 100644 --- a/Source/santa-driver/SantaCache.h +++ b/Source/santa-driver/SantaCache.h @@ -101,12 +101,19 @@ template 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 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 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 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 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 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 diff --git a/Source/santa-driver/SantaDecisionManager.cc b/Source/santa-driver/SantaDecisionManager.cc index d8445890..5a0bb8da 100644 --- a/Source/santa-driver/SantaDecisionManager.cc +++ b/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(this)); + KAUTH_SCOPE_FILEOP, fileop_scope_callback, + reinterpret_cast(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* SantaDecisionManager::CacheForIdentifier(const uint64_t identifier) { - return (identifier >> 32 == root_fsid_) ? root_decision_cache_ : non_root_decision_cache_; +SantaCache* 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; diff --git a/Tests/LogicTests/SantaCacheTest.mm b/Tests/LogicTests/SantaCacheTest.mm index d008695a..22311f5f 100644 --- a/Tests/LogicTests/SantaCacheTest.mm +++ b/Tests/LogicTests/SantaCacheTest.mm @@ -150,4 +150,27 @@ delete sut; } +- (void)testCompareAndSwap { + auto sut = new SantaCache(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