santad: Wait for driver appearance using IOKit notifications. (#278)

Continue loading without driver, report status in santactl.

Conf/com.google.santad.plist

@@ -20,7 +20,5 @@
 				<true />
 				<key>ProcessType</key>
 				<string>Interactive</string>
-				<key>ThrottleInterval</key>
-				<integer>1</integer>
 </dict>
 </plist>

Docs/deployment/configuration.md

@@ -37,6 +37,7 @@ Two configuration methods can be used to control Santa: a local configuration pr
 | MachineIDKey                  | String     | The key to use on MachineIDPlist.        |
 | EventLogType                  | String     | Defines how event logs are stored. Options are 1) syslog: Sent to ASL or ULS (if built with the 10.12 SDK or later). 2) filelog: Sent to a file on disk. Use EventLogPath to specify a path. Defaults to filelog      |
 | EventLogPath                  | String     | If EventLogType is set to filelog, EventLogPath will provide the path to save logs. Defaults to /var/db/santa/santa.log. If you change this value ensure you also update com.google.santa.newsyslog.conf with the new path.        |
+| EnableMachineIDDecoration     | Bool       | If YES, this appends the MachineID to the end of each log line. Defaults to NO.       |
 
 *overridable by the sync server: run `santactl status` to check the current running config
 

Podfile.lock

@@ -22,6 +22,16 @@ DEPENDENCIES:
   - MOLXPCConnection
   - OCMock
 
+SPEC REPOS:
+  https://github.com/cocoapods/specs.git:
+    - FMDB
+    - MOLAuthenticatingURLSession
+    - MOLCertificate
+    - MOLCodesignChecker
+    - MOLFCMClient
+    - MOLXPCConnection
+    - OCMock
+
 SPEC CHECKSUMS:
   FMDB: 6198a90e7b6900cfc046e6bc0ef6ebb7be9236aa
   MOLAuthenticatingURLSession: c238aa1c9a7b1077eb39a6f40204bfe76a7d204e
@@ -33,4 +43,4 @@ SPEC CHECKSUMS:
 
 PODFILE CHECKSUM: ddca043a7ace9ec600c108621c56d13a50d17236
 
-COCOAPODS: 1.4.0
+COCOAPODS: 1.5.3

README.md

@@ -43,7 +43,7 @@ correctly, but a rule for a binary's fingerprint will override a decision for a
 certificate; i.e. you can whitelist a certificate while blacklisting a binary
 signed with that certificate, or vice-versa.
 
-* Path-based rules (via NSRegularExpression/ICU): This allows a similar feature to that found in Managed Client (the precursor to configuration profiles, which used the same implementation mechanism), Application Launch Restrictions via the mcxalr binary. This implementation carries the added benefit of being configurable via regex, and not relying on LaunchServices. As detailed in the wiki, when evaluating rules this holds the lowest precendence.
+* Path-based rules (via NSRegularExpression/ICU): This allows a similar feature to that found in Managed Client (the precursor to configuration profiles, which used the same implementation mechanism), Application Launch Restrictions via the mcxalr binary. This implementation carries the added benefit of being configurable via regex, and not relying on LaunchServices. As detailed in the wiki, when evaluating rules this holds the lowest precedence.
 
 * Failsafe cert rules: You cannot put in a deny rule that would block the certificate used to sign launchd, a.k.a. pid 1, and therefore all components used in macOS. The binaries in every OS update (and in some cases entire new versions) are therefore auto-whitelisted. This does not affect binaries from Apple's App Store, which use various certs that change regularly for common apps. Likewise, you cannot blacklist Santa itself, and Santa uses a distinct separate cert than other Google apps.
 

Rakefile

@@ -155,20 +155,8 @@ namespace :tests do
 
   desc "Tests: Kernel"
   task :kernel do
-    Rake::Task['unload'].invoke()
-    Rake::Task['install:debug'].invoke()
-    Rake::Task['load_kext'].invoke
-    FileUtils.mkdir_p("/tmp/santa_kerneltests_tmp")
-    begin
-      puts "\033[?25l\033[12h"  # hide cursor
-      puts "Running kernel tests"
-      system "cd /tmp/santa_kerneltests_tmp && sudo #{xcodebuilddir}/Debug/KernelTests"
-    rescue Exception
-    ensure
-      puts "\033[?25h\033[12l\n\n"  # unhide cursor
-      FileUtils.rm_rf("/tmp/santa_kerneltests_tmp")
-      Rake::Task['unload_kext'].execute
-    end
+    Rake::Task['build:debug'].invoke()
+    system "sudo #{xcodebuilddir}/Debug/KernelTests"
   end
 end
 

Santa.xcodeproj/project.pbxproj

@@ -32,6 +32,7 @@
 		0D10BE861A0AABD600C0C944 /* SNTDropRootPrivs.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D10BE851A0AABD600C0C944 /* SNTDropRootPrivs.m */; };
 		0D10BE871A0AABD600C0C944 /* SNTDropRootPrivs.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D10BE851A0AABD600C0C944 /* SNTDropRootPrivs.m */; };
 		0D10BE891A0AAF6700C0C944 /* SNTDropRootPrivs.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D10BE851A0AABD600C0C944 /* SNTDropRootPrivs.m */; };
+		0D10D18420C19445008251ED /* SNTCommandCacheHistogram.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D10D18320C19445008251ED /* SNTCommandCacheHistogram.m */; };
 		0D1B477019A53419008CADD3 /* SNTAboutWindowController.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D1B476E19A53419008CADD3 /* SNTAboutWindowController.m */; };
 		0D1B477119A53419008CADD3 /* AboutWindow.xib in Resources */ = {isa = PBXBuildFile; fileRef = 0D1B476F19A53419008CADD3 /* AboutWindow.xib */; };
 		0D202D191CDD2EE500A88F16 /* SNTCommandSyncTest.m in Sources */ = {isa = PBXBuildFile; fileRef = 0D202D181CDD2EE500A88F16 /* SNTCommandSyncTest.m */; };
@@ -286,6 +287,7 @@
 		0D0A1EC5191AB9B000B8450F /* SNTCommandSyncPostflight.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = SNTCommandSyncPostflight.m; sourceTree = "<group>"; };
 		0D10BE851A0AABD600C0C944 /* SNTDropRootPrivs.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = SNTDropRootPrivs.m; sourceTree = "<group>"; };
 		0D10BE881A0AAC2100C0C944 /* SNTDropRootPrivs.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = SNTDropRootPrivs.h; sourceTree = "<group>"; };
+		0D10D18320C19445008251ED /* SNTCommandCacheHistogram.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = SNTCommandCacheHistogram.m; sourceTree = "<group>"; };
 		0D1B476D19A53419008CADD3 /* SNTAboutWindowController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = SNTAboutWindowController.h; sourceTree = "<group>"; };
 		0D1B476E19A53419008CADD3 /* SNTAboutWindowController.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = SNTAboutWindowController.m; sourceTree = "<group>"; };
 		0D1B476F19A53419008CADD3 /* AboutWindow.xib */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = file.xib; path = AboutWindow.xib; sourceTree = "<group>"; };
@@ -801,6 +803,7 @@
 			isa = PBXGroup;
 			children = (
 				C7DA62F81E241A02009BDF2C /* SNTCommandBundleInfo.m */,
+				0D10D18320C19445008251ED /* SNTCommandCacheHistogram.m */,
 				C76614EB1D142D3C00D150C1 /* SNTCommandCheckCache.m */,
 				0DCD5FBE1909D64A006B445C /* SNTCommandFileInfo.m */,
 				0DE4C8A518FF3B1700466D04 /* SNTCommandFlushCache.m */,
@@ -929,8 +932,6 @@
 				0D260DA818B68E12002A0B55 /* Sources */,
 				0D260DA918B68E12002A0B55 /* Frameworks */,
 				0D260DAA18B68E12002A0B55 /* Resources */,
-				D49A3AB950AFD99741E9AF89 /* [CP] Embed Pods Frameworks */,
-				23869BA352E2C86DEFE62819 /* [CP] Copy Pods Resources */,
 			);
 			buildRules = (
 			);
@@ -949,7 +950,6 @@
 				0DD98E671A5DD02000A754C6 /* Update Version Info */,
 				0D35BD9A18FD71CE00921A21 /* Sources */,
 				0D35BD9B18FD71CE00921A21 /* Frameworks */,
-				2737FE8516A33567D3449943 /* [CP] Copy Pods Resources */,
 			);
 			buildRules = (
 			);
@@ -969,8 +969,6 @@
 				0D385DB2180DE4A900418BC6 /* Sources */,
 				0D385DB3180DE4A900418BC6 /* Frameworks */,
 				0D385DB4180DE4A900418BC6 /* Resources */,
-				BA20035148DDEF5808B2C7EF /* [CP] Embed Pods Frameworks */,
-				DDE76075391100F3BCE3634E /* [CP] Copy Pods Resources */,
 			);
 			buildRules = (
 			);
@@ -1012,7 +1010,6 @@
 				0DD98E661A5DCED300A754C6 /* Update Version Info */,
 				0D9A7F391759330400035EB5 /* Sources */,
 				0D9A7F3A1759330400035EB5 /* Frameworks */,
-				531BBB5BA4BF953D628D29AB /* [CP] Copy Pods Resources */,
 			);
 			buildRules = (
 			);
@@ -1033,7 +1030,6 @@
 				C78227501E1C3C58006EB2D6 /* Sources */,
 				C78227511E1C3C58006EB2D6 /* Frameworks */,
 				C78227521E1C3C58006EB2D6 /* Resources */,
-				FCF4BDAA180BD63C32AC85DE /* [CP] Copy Pods Resources */,
 			);
 			buildRules = (
 			);
@@ -1253,51 +1249,6 @@
 			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
 			showEnvVarsInLog = 0;
 		};
-		23869BA352E2C86DEFE62819 /* [CP] Copy Pods Resources */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Copy Pods Resources";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-LogicTests/Pods-LogicTests-resources.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
-		2737FE8516A33567D3449943 /* [CP] Copy Pods Resources */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Copy Pods Resources";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-santactl/Pods-santactl-resources.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
-		531BBB5BA4BF953D628D29AB /* [CP] Copy Pods Resources */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Copy Pods Resources";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-santad/Pods-santad-resources.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
 		9BF8830B1029605A497F13D5 /* [CP] Check Pods Manifest.lock */ = {
 			isa = PBXShellScriptBuildPhase;
 			buildActionMask = 2147483647;
@@ -1334,21 +1285,6 @@
 			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
 			showEnvVarsInLog = 0;
 		};
-		BA20035148DDEF5808B2C7EF /* [CP] Embed Pods Frameworks */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Embed Pods Frameworks";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-Santa/Pods-Santa-frameworks.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
 		C7F95F971E23F9E7007A6BF5 /* Update bundle version */ = {
 			isa = PBXShellScriptBuildPhase;
 			buildActionMask = 2147483647;
@@ -1384,51 +1320,6 @@
 			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
 			showEnvVarsInLog = 0;
 		};
-		D49A3AB950AFD99741E9AF89 /* [CP] Embed Pods Frameworks */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Embed Pods Frameworks";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-LogicTests/Pods-LogicTests-frameworks.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
-		DDE76075391100F3BCE3634E /* [CP] Copy Pods Resources */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Copy Pods Resources";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-Santa/Pods-Santa-resources.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
-		FCF4BDAA180BD63C32AC85DE /* [CP] Copy Pods Resources */ = {
-			isa = PBXShellScriptBuildPhase;
-			buildActionMask = 2147483647;
-			files = (
-			);
-			inputPaths = (
-			);
-			name = "[CP] Copy Pods Resources";
-			outputPaths = (
-			);
-			runOnlyForDeploymentPostprocessing = 0;
-			shellPath = /bin/sh;
-			shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-santad-santabs/Pods-santad-santabs-resources.sh\"\n";
-			showEnvVarsInLog = 0;
-		};
 /* End PBXShellScriptBuildPhase section */
 
 /* Begin PBXSourcesBuildPhase section */
@@ -1505,6 +1396,7 @@
 				C776A1071DEE160500A56616 /* SNTCommandSyncManager.m in Sources */,
 				0DCD605C19117A90006B445C /* SNTCommandSyncPreflight.m in Sources */,
 				0D41640519197AD7006A356A /* SNTCommandSyncEventUpload.m in Sources */,
+				0D10D18420C19445008251ED /* SNTCommandCacheHistogram.m in Sources */,
 				0D42D2B919D2042900955F08 /* SNTConfigurator.m in Sources */,
 				0DF395641AB76A7900CBC520 /* NSData+Zlib.m in Sources */,
 				C7FB56F71DBFB480004E14EF /* SNTXPCSyncdInterface.m in Sources */,
@@ -1692,6 +1584,7 @@
 				GCC_WARN_UNUSED_FUNCTION = YES;
 				GCC_WARN_UNUSED_VARIABLE = YES;
 				INSTALL_PATH = "";
+				OTHER_CODE_SIGN_FLAGS = "";
 				OTHER_CPLUSPLUSFLAGS = (
 					"$(OTHER_CFLAGS)",
 					"-fcxx-modules",
@@ -1726,6 +1619,7 @@
 				GCC_WARN_UNUSED_FUNCTION = YES;
 				GCC_WARN_UNUSED_VARIABLE = YES;
 				INSTALL_PATH = "";
+				OTHER_CODE_SIGN_FLAGS = "";
 				OTHER_CPLUSPLUSFLAGS = (
 					"$(OTHER_CFLAGS)",
 					"-fcxx-modules",

Santa.xcodeproj/xcshareddata/xcschemes/KernelTests.xcscheme

@@ -46,6 +46,7 @@
       buildConfiguration = "Debug"
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      debugAsWhichUser = "root"
       language = ""
       launchStyle = "0"
       useCustomWorkingDirectory = "NO"

Source/common/SNTConfigurator.h

@@ -98,6 +98,13 @@
 ///
 @property(readonly, nonatomic) NSString *eventLogPath;
 
+///
+/// Enabling this appends the Santa machine ID to the end of each log line. If nothing
+/// has been overriden, this is the host's UUID.
+/// Defaults to NO.
+///
+@property(readonly, nonatomic) BOOL enableMachineIDDecoration;
+
 #pragma mark - GUI Settings
 
 ///

Source/common/SNTConfigurator.m

@@ -72,6 +72,8 @@ static NSString *const kFileChangesRegexKey = @"FileChangesRegex";
 static NSString *const kEventLogType = @"EventLogType";
 static NSString *const kEventLogPath = @"EventLogPath";
 
+static NSString *const kEnableMachineIDDecoration = @"EnableMachineIDDecoration";
+
 // The keys managed by a sync server or mobileconfig.
 static NSString *const kClientModeKey = @"ClientMode";
 static NSString *const kWhitelistRegexKey = @"WhitelistRegex";
@@ -126,6 +128,7 @@ static NSString *const kSyncCleanRequired = @"SyncCleanRequired";
       kMachineIDPlistKeyKey : string,
       kEventLogType : string,
       kEventLogPath : string,
+      kEnableMachineIDDecoration : number,
     };
     _defaults = [NSUserDefaults standardUserDefaults];
     [_defaults addSuiteNamed:@"com.google.santa"];
@@ -280,6 +283,10 @@ static NSString *const kSyncCleanRequired = @"SyncCleanRequired";
   return [self configStateSet];
 }
 
++ (NSSet *)keyPathsForValuesAffectingEnableMachineIDDecoration {
+  return [self configStateSet];
+}
+
 #pragma mark Public Interface
 
 - (SNTClientMode)clientMode {
@@ -451,6 +458,11 @@ static NSString *const kSyncCleanRequired = @"SyncCleanRequired";
   return self.configState[kEventLogPath] ?: @"/var/db/santa/santa.log";
 }
 
+- (BOOL)enableMachineIDDecoration {
+  NSNumber *number = self.configState[kEnableMachineIDDecoration];
+  return number ? [number boolValue] : NO;
+}
+
 #pragma mark Private
 
 ///

Source/common/SNTKernelCommon.h

@@ -21,9 +21,6 @@
 #ifndef SANTA__COMMON__KERNELCOMMON_H
 #define SANTA__COMMON__KERNELCOMMON_H
 
-// Defines the lengths of paths and Vnode IDs passed around.
-#define MAX_VNODE_ID_STR 21  // digits in UINT64_MAX + 1 for NULL-terminator
-
 // Defines the name of the userclient class and the driver bundle ID.
 #define USERCLIENT_CLASS "com_google_SantaDriver"
 #define USERCLIENT_ID "com.google.santa-driver"
@@ -41,6 +38,7 @@ enum SantaDriverMethods {
   kSantaUserClientClearCache,
   kSantaUserClientCacheCount,
   kSantaUserClientCheckCache,
+  kSantaUserClientCacheBucketCount,
 
   // Any methods supported by the driver should be added above this line to
   // ensure this remains the count of methods.
@@ -82,10 +80,27 @@ typedef enum {
 #define RESPONSE_VALID(x) \
   (x == ACTION_RESPOND_ALLOW || x == ACTION_RESPOND_DENY)
 
+// Struct to manage vnode IDs
+typedef struct santa_vnode_id_t {
+  uint64_t fsid;
+  uint64_t fileid;
+
+#ifdef __cplusplus
+  bool operator==(const santa_vnode_id_t& rhs) const {
+    return fsid == rhs.fsid && fileid == rhs.fileid;
+  }
+  // This _must not_ be used for anything security-sensitive. It exists solely to make
+  // the msleep/wakeup calls easier.
+  uint64_t unsafe_simple_id() const {
+    return (((uint64_t)fsid << 32) | fileid);
+  }
+#endif
+} santa_vnode_id_t;
+
 // Message struct that is sent down the IODataQueue.
 typedef struct {
   santa_action_t action;
-  uint64_t vnode_id;
+  santa_vnode_id_t vnode_id;
   uid_t uid;
   gid_t gid;
   pid_t pid;
@@ -99,4 +114,10 @@ typedef struct {
   char pname[MAXPATHLEN];
 } santa_message_t;
 
+// Used for the kSantaUserClientCacheBucketCount request.
+typedef struct {
+  uint16_t per_bucket[1024];
+  uint64_t start;
+} santa_bucket_count_t;
+
 #endif  // SANTA__COMMON__KERNELCOMMON_H

Source/common/SNTXPCControlInterface.h

@@ -33,9 +33,11 @@
 ///
 ///  Kernel ops
 ///
-- (void)cacheCounts:(void (^)(uint64_t rootCache, uint64_t nonRootCache))reply;
+- (void)cacheCounts:(void (^)(uint64_t count))reply;
 - (void)flushCache:(void (^)(BOOL))reply;
-- (void)checkCacheForVnodeID:(uint64_t)vnodeID withReply:(void (^)(santa_action_t))reply;
+- (void)cacheBucketCount:(void (^)(NSArray *))reply;
+- (void)checkCacheForVnodeID:(santa_vnode_id_t)vnodeID withReply:(void (^)(santa_action_t))reply;
+- (void)driverConnectionEstablished:(void (^)(BOOL))reply;
 
 ///
 ///  Database ops

Source/santa-driver/SantaCache.h

@@ -40,9 +40,21 @@
 #define OSDecrementAtomic(addr) OSAtomicDecrement64((volatile int64_t *)addr)
 #endif // KERNEL
 
+/**
+  A type to specialize to help SantaCache with its hashing.
+
+  The default works for numeric types with a multiplicative hash
+  using a prime near to the golden ratio, per Knuth.
+*/
+template<typename T> uint64_t SantaCacheHasher(T const& t) {
+  return t * 11400714819323198549UL;
+};
+
 /**
   A somewhat simple, concurrent linked-list hash table intended for use in IOKit kernel extensions.
-  Maps 64-bit unsigned integer keys to values.
+
+  The type used for keys must overload the == operator and a specialization of
+  SantaCacheHasher must exist for it.
 
   Enforces a maximum size by clearing all entries if a new value
   is added that would go over the maximum size declared at creation.
@@ -50,7 +62,7 @@
   The number of buckets is calculated as `maximum_size` / `per_bucket`
   rounded up to the next power of 2. Locking is done per-bucket.
 */
-template<class T> class SantaCache {
+template<typename KeyT, typename ValueT> class SantaCache {
  public:
   /**
     Initialize a newly created cache.
@@ -65,8 +77,7 @@ template<class T> class SantaCache {
     if (unlikely(per_bucket < 1)) per_bucket = 1;
     if (unlikely(per_bucket > 64)) per_bucket = 64;
     max_size_ = maximum_size;
-    bucket_count_ = 1 << (32 - __builtin_clz(
-        ((uint32_t)max_size_ / per_bucket) - 1));
+    bucket_count_ = 1 << (32 - __builtin_clz(((uint32_t)max_size_ / per_bucket) - 1));
     buckets_ = (struct bucket *)IOMalloc(bucket_count_ * sizeof(struct bucket));
     bzero(buckets_, bucket_count_ * sizeof(struct bucket));
   }
@@ -82,13 +93,13 @@ template<class T> class SantaCache {
   /**
     Get an element from the cache. Returns zero_ if item doesn't exist.
   */
-  T get(uint64_t key) {
+  ValueT get(KeyT key) {
     struct bucket *bucket = &buckets_[hash(key)];
     lock(bucket);
     struct entry *entry = (struct entry *)((uintptr_t)bucket->head - 1);
     while (entry != nullptr) {
       if (entry->key == key) {
-        T val = entry->value;
+        ValueT val = entry->value;
         unlock(bucket);
         return val;
       }
@@ -102,103 +113,39 @@ 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.
+        empty the cache before inserting the new value.
 
-    @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.
+    @param key, The key.
+    @param value, The value with parameterized type.
 
-    @return the previous value (which may be zero_)
+    @return true if the value was set.
   */
-  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);
-    struct entry *previous_entry = nullptr;
-    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_) {
-          if (previous_entry != nullptr) {
-            previous_entry->next = entry->next;
-          } else {
-            bucket->head = (struct entry *)((uintptr_t)entry->next + 1);
-          }
-          IOFreeAligned(entry, sizeof(struct entry));
-          OSDecrementAtomic(&count_);
-        }
-
-        unlock(bucket);
-        return existing_value;
-      }
-      previous_entry = entry;
-      entry = entry->next;
-    }
-
-    // If value is zero_, we're clearing but there's nothing to clear
-    // 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.
-    if (count_ + 1 > max_size_) {
-      unlock(bucket);
-      lock(&clear_bucket_);
-      // Check again in case clear has already run while waiting for lock
-      if (count_ + 1 > max_size_) {
-        clear();
-      }
-      lock(bucket);
-      unlock(&clear_bucket_);
-    }
-
-    // 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);
-    bucket->head = (struct entry *)((uintptr_t)new_entry + 1);
-    OSIncrementAtomic(&count_);
-
-    unlock(bucket);
-    return zero_;
-  }
-
-  /**
-    Overload to allow setting without providing a previous value
-  */
-  T set(uint64_t key, T value) {
+  bool set(KeyT key, ValueT value) {
     return set(key, value, {}, false);
   }
 
   /**
-    Overload to allow setting while providing a previous value
+    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.
+
+    @param key, The key.
+    @param value, The value with parameterized type.
+    @param previous_value, the new value will only be set if this
+        parameter is equal to the existing value in the cache.
+        This allows set to become a CAS operation.
+
+    @return true if the value was set
   */
-  T set(uint64_t key, T value, T previous_value) {
+  bool set(KeyT key, ValueT value, ValueT previous_value) {
     return set(key, value, previous_value, true);
   }
 
   /**
     An alias for `set(key, zero_)`
   */
-  inline void remove(uint64_t key) {
+  inline void remove(KeyT key) {
     set(key, zero_);
   }
 
@@ -238,10 +185,42 @@ template<class T> class SantaCache {
     return count_;
   }
 
+  /**
+    Fill in the per_bucket_counts array with the number of entries in each bucket.
+
+    The per_buckets_count array will contain the per-bucket counts, up to the number
+    in array_size. The start_bucket parameter will determine which bucket to start off
+    with and upon return will contain either 0 if no buckets are remaining or the next
+    bucket to begin with when called again.
+  */
+  void bucket_counts(uint16_t *per_bucket_counts, uint16_t *array_size, uint64_t *start_bucket) {
+    if (per_bucket_counts == nullptr || array_size == nullptr || start_bucket == nullptr) return;
+
+    uint64_t start = *start_bucket;
+    uint16_t size = *array_size;
+    if (start + size > bucket_count_) size = bucket_count_ - start;
+
+    for (uint16_t i = 0; i < size; ++i) {
+      uint16_t count = 0;
+      struct bucket *bucket = &buckets_[start++];
+      lock(bucket);
+      struct entry *entry = (struct entry *)((uintptr_t)bucket->head - 1);
+      while (entry != nullptr) {
+        if (entry->value != zero_) ++count;
+        entry = entry->next;
+      }
+      unlock(bucket);
+      per_bucket_counts[i] = count;
+    }
+
+    *array_size = size;
+    *start_bucket = (start >= bucket_count_) ? 0 : start;
+  }
+
  private:
   struct entry {
-    uint64_t key;
-    T value;
+    KeyT key;
+    ValueT value;
     struct entry *next;
   };
 
@@ -251,6 +230,88 @@ template<class T> class SantaCache {
     struct entry *head;
   };
 
+  /**
+    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.
+
+    @param key, The key
+    @param value, The value with parameterized type
+    @param previous_value, If has_prev_value is true, the new value will only
+        be set if this parameter is equal to the existing value in the cache.
+        This allows set to become a CAS operation.
+    @param has_prev_value, Pass true if previous_value should be used.
+
+    @return true if the entry was set, false if it was not
+  */
+  bool set(KeyT key, ValueT value, ValueT previous_value, bool has_prev_value) {
+    struct bucket *bucket = &buckets_[hash(key)];
+    lock(bucket);
+    struct entry *entry = (struct entry *)((uintptr_t)bucket->head - 1);
+    struct entry *previous_entry = nullptr;
+    while (entry != nullptr) {
+      if (entry->key == key) {
+        ValueT existing_value = entry->value;
+
+        if (has_prev_value && previous_value != existing_value) {
+          unlock(bucket);
+          return false;
+        }
+
+        entry->value = value;
+
+        if (value == zero_) {
+          if (previous_entry != nullptr) {
+            previous_entry->next = entry->next;
+          } else {
+            bucket->head = (struct entry *)((uintptr_t)entry->next + 1);
+          }
+          IOFreeAligned(entry, sizeof(struct entry));
+          OSDecrementAtomic(&count_);
+        }
+
+        unlock(bucket);
+        return true;
+      }
+      previous_entry = entry;
+      entry = entry->next;
+    }
+
+    // If value is zero_, we're clearing but there's nothing to clear
+    // 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 false;
+    }
+
+    // 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_);
+      // Check again in case clear has already run while waiting for lock
+      if (count_ + 1 > max_size_) {
+        clear();
+      }
+      lock(bucket);
+      unlock(&clear_bucket_);
+    }
+
+    // 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);
+    bucket->head = (struct entry *)((uintptr_t)new_entry + 1);
+    OSIncrementAtomic(&count_);
+
+    unlock(bucket);
+    return true;
+  }
+
   /**
     Lock a bucket. Spins until the lock is acquired.
   */
@@ -277,7 +338,7 @@ template<class T> class SantaCache {
   /**
     Holder for a 'zero' entry for the current type
   */
-  const T zero_ = {};
+  const ValueT zero_ = {};
 
   /**
     Special bucket used when automatically clearing due to size
@@ -288,13 +349,9 @@ template<class T> class SantaCache {
 
   /**
     Hash a key to determine which bucket it belongs in.
-
-    Multiplicative hash using a prime near to the golden ratio, per Knuth.
-    This seems to have good bucket distribution generally and for the range of
-    values we expect to see.
   */
-  inline uint64_t hash(uint64_t input) const {
-    return (input * 11400714819323198549ul) % bucket_count_;
+  inline uint64_t hash(KeyT input) const {
+    return SantaCacheHasher<KeyT>(input) % bucket_count_;
   }
 };
 

Source/santa-driver/SantaDecisionManager.cc

@@ -19,6 +19,10 @@ OSDefineMetaClassAndStructors(SantaDecisionManager, OSObject);
 
 #pragma mark Object Lifecycle
 
+template<> uint64_t SantaCacheHasher<santa_vnode_id_t>(santa_vnode_id_t const& s) {
+  return (SantaCacheHasher<uint64_t>(s.fsid) << 1) ^ SantaCacheHasher<uint64_t>(s.fileid);
+}
+
 bool SantaDecisionManager::init() {
   if (!super::init()) return false;
 
@@ -29,9 +33,8 @@ bool SantaDecisionManager::init() {
   decision_dataqueue_lock_ = lck_mtx_alloc_init(sdm_lock_grp_, sdm_lock_attr_);
   log_dataqueue_lock_ = lck_mtx_alloc_init(sdm_lock_grp_, sdm_lock_attr_);
 
-  root_decision_cache_ = new SantaCache<uint64_t>(5000, 2);
-  non_root_decision_cache_ = new SantaCache<uint64_t>(500, 2);
-  vnode_pid_map_ = new SantaCache<uint64_t>(2000, 5);
+  decision_cache_ = new SantaCache<santa_vnode_id_t, uint64_t>(10000, 2);
+  vnode_pid_map_ = new SantaCache<santa_vnode_id_t, uint64_t>(2000, 5);
 
   decision_dataqueue_ = IOSharedDataQueue::withEntries(
       kMaxDecisionQueueEvents, sizeof(santa_message_t));
@@ -42,17 +45,12 @@ bool SantaDecisionManager::init() {
   if (!log_dataqueue_) return kIOReturnNoMemory;
 
   client_pid_ = 0;
-  root_fsid_ = 0;
-
-  ts_ = { .tv_sec = kRequestLoopSleepMilliseconds / 1000,
-          .tv_nsec = kRequestLoopSleepMilliseconds % 1000 * 1000000 };
 
   return true;
 }
 
 void SantaDecisionManager::free() {
-  delete root_decision_cache_;
-  delete non_root_decision_cache_;
+  delete decision_cache_;
   delete vnode_pid_map_;
 
   if (decision_dataqueue_lock_) {
@@ -93,17 +91,6 @@ void SantaDecisionManager::ConnectClient(pid_t pid) {
 
   client_pid_ = pid;
 
-  // Determine root fsid
-  vfs_context_t ctx = vfs_context_create(NULL);
-  if (ctx) {
-    vnode_t root = vfs_rootvnode();
-    if (root) {
-      root_fsid_ = GetVnodeIDForVnode(ctx, root) >> 32;
-      vnode_put(root);
-    }
-    vfs_context_rele(ctx);
-  }
-
   // Any decisions made while the daemon wasn't
   // connected should be cleared
   ClearCache();
@@ -112,8 +99,8 @@ void SantaDecisionManager::ConnectClient(pid_t pid) {
   failed_log_queue_requests_ = 0;
 }
 
-void SantaDecisionManager::DisconnectClient(bool itDied) {
-  if (client_pid_ < 1) return;
+void SantaDecisionManager::DisconnectClient(bool itDied, pid_t pid) {
+  if (client_pid_ == 0 || (pid > 0 && pid != client_pid_)) return;
   client_pid_ = 0;
 
   // Ask santad to shutdown, in case it's running.
@@ -210,74 +197,58 @@ kern_return_t SantaDecisionManager::StopListener() {
 
 #pragma mark Cache Management
 
-/**
-  Return the correct cache for a given identifier.
-
-  @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_;
-}
-
 void SantaDecisionManager::AddToCache(
-    uint64_t identifier, santa_action_t decision, uint64_t microsecs) {
-  // Decision is stored in upper 8 bits, timestamp in remaining 56.
-  uint64_t val = ((uint64_t)decision << 56) | (microsecs & 0xFFFFFFFFFFFFFF);
-
-  auto decision_cache = CacheForIdentifier(identifier);
-
+    santa_vnode_id_t identifier, santa_action_t decision, uint64_t microsecs) {
   switch (decision) {
     case ACTION_REQUEST_BINARY:
-      decision_cache->set(identifier, val, 0);
+      decision_cache_->set(identifier, (uint64_t)ACTION_REQUEST_BINARY << 56, 0);
       break;
     case ACTION_RESPOND_ACK:
-      decision_cache->set(identifier, val, ((uint64_t)ACTION_REQUEST_BINARY << 56));
+      decision_cache_->set(identifier, (uint64_t)ACTION_RESPOND_ACK << 56,
+                          ((uint64_t)ACTION_REQUEST_BINARY << 56));
       break;
     case ACTION_RESPOND_ALLOW:
-    case ACTION_RESPOND_DENY:
-      // TODO(bur): Avoid calling set() twice, finding and locking buckets is fast, but not free.
-      if (decision_cache->set(identifier, val, ((uint64_t)ACTION_REQUEST_BINARY << 56))) {
-        decision_cache->set(identifier, val, ((uint64_t)ACTION_RESPOND_ACK << 56));
+    case ACTION_RESPOND_DENY: {
+      // Decision is stored in upper 8 bits, timestamp in remaining 56.
+      uint64_t val = ((uint64_t)decision << 56) | (microsecs & 0xFFFFFFFFFFFFFF);
+      if (!decision_cache_->set(identifier, val, ((uint64_t)ACTION_REQUEST_BINARY << 56))) {
+        decision_cache_->set(identifier, val, ((uint64_t)ACTION_RESPOND_ACK << 56));
       }
       break;
+    }
     default:
       break;
   }
 
-  wakeup((void *)identifier);
+  wakeup((void *)identifier.unsafe_simple_id());
 }
 
-void SantaDecisionManager::RemoveFromCache(uint64_t identifier) {
-  CacheForIdentifier(identifier)->remove(identifier);
-  if (unlikely(!identifier)) return;
-  wakeup((void *)identifier);
+void SantaDecisionManager::RemoveFromCache(santa_vnode_id_t identifier) {
+  if (unlikely(identifier.fsid == 0 && identifier.fileid == 0)) return;
+  decision_cache_->remove(identifier);
+  wakeup((void *)identifier.unsafe_simple_id());
 }
 
-uint64_t SantaDecisionManager::RootCacheCount() const {
-  return root_decision_cache_->count();
+uint64_t SantaDecisionManager::CacheCount() const {
+  return decision_cache_->count();
 }
 
-uint64_t SantaDecisionManager::NonRootCacheCount() const {
-  return non_root_decision_cache_->count();
+void SantaDecisionManager::ClearCache() {
+  decision_cache_->clear();
 }
 
-void SantaDecisionManager::ClearCache(bool non_root_only) {
-  if (!non_root_only) root_decision_cache_->clear();
-  non_root_decision_cache_->clear();
+void SantaDecisionManager::CacheBucketCount(
+    uint16_t *per_bucket_counts, uint16_t *array_size, uint64_t *start_bucket) {
+  decision_cache_->bucket_counts(per_bucket_counts, array_size, start_bucket);
 }
 
 #pragma mark Decision Fetching
 
-santa_action_t SantaDecisionManager::GetFromCache(uint64_t identifier) {
+santa_action_t SantaDecisionManager::GetFromCache(santa_vnode_id_t identifier) {
   auto result = ACTION_UNSET;
   uint64_t decision_time = 0;
 
-  auto decision_cache = CacheForIdentifier(identifier);
-
-  uint64_t cache_val = decision_cache->get(identifier);
+  uint64_t cache_val = decision_cache_->get(identifier);
   if (cache_val == 0) return result;
 
   // Decision is stored in upper 8 bits, timestamp in remaining 56.
@@ -288,7 +259,7 @@ santa_action_t SantaDecisionManager::GetFromCache(uint64_t identifier) {
     if (result == ACTION_RESPOND_DENY) {
       auto expiry_time = decision_time + (kMaxDenyCacheTimeMilliseconds * 1000);
       if (expiry_time < GetCurrentUptime()) {
-        decision_cache->remove(identifier);
+        decision_cache_->remove(identifier);
         return ACTION_UNSET;
       }
     }
@@ -298,7 +269,7 @@ santa_action_t SantaDecisionManager::GetFromCache(uint64_t identifier) {
 }
 
 santa_action_t SantaDecisionManager::GetFromDaemon(
-    santa_message_t *message, uint64_t identifier) {
+    santa_message_t *message, santa_vnode_id_t identifier) {
   auto return_action = ACTION_UNSET;
 
 #ifdef DEBUG
@@ -326,7 +297,7 @@ santa_action_t SantaDecisionManager::GetFromDaemon(
     // request, indicated with ACTION_RESPOND_ACK.
     auto cache_check_count = 0;
     do {
-      msleep((void *)message->vnode_id, NULL, 0, "", &ts_);
+      msleep((void *)message->vnode_id.unsafe_simple_id(), NULL, 0, "", &ts_);
       return_action = GetFromCache(identifier);
     } while (ClientConnected() &&
              ((return_action == ACTION_REQUEST_BINARY && ++cache_check_count < kRequestCacheChecks)
@@ -353,7 +324,7 @@ santa_action_t SantaDecisionManager::GetFromDaemon(
 santa_action_t SantaDecisionManager::FetchDecision(
     const kauth_cred_t cred,
     const vnode_t vp,
-    const uint64_t vnode_id) {
+    const santa_vnode_id_t vnode_id) {
   while (true) {
     if (!ClientConnected()) return ACTION_RESPOND_ALLOW;
 
@@ -368,7 +339,7 @@ santa_action_t SantaDecisionManager::FetchDecision(
     } else if (return_action == ACTION_REQUEST_BINARY || return_action == ACTION_RESPOND_ACK) {
       // 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_);
+      msleep((void *)vnode_id.unsafe_simple_id(), NULL, 0, "", &ts_);
     } else {
       break;
     }
@@ -403,7 +374,6 @@ bool SantaDecisionManager::PostToDecisionQueue(santa_message_t *message) {
       LOGE("Failed to queue more than %d decision requests, killing daemon",
            kMaxDecisionQueueFailures);
       proc_signal(client_pid_, SIGKILL);
-      client_pid_ = 0;
     }
   }
   lck_mtx_unlock(decision_dataqueue_lock_);
@@ -417,10 +387,6 @@ bool SantaDecisionManager::PostToLogQueue(santa_message_t *message) {
     if (failed_log_queue_requests_++ == 0) {
       LOGW("Dropping log queue messages");
     }
-    // If enqueue failed, pop an item off the queue and try again.
-    uint32_t dataSize = 0;
-    log_dataqueue_->dequeue(0, &dataSize);
-    kr = log_dataqueue_->enqueue(message, sizeof(santa_message_t));
   } else {
     if (failed_log_queue_requests_ > 0) {
       failed_log_queue_requests_--;
@@ -448,24 +414,11 @@ int SantaDecisionManager::VnodeCallback(const kauth_cred_t cred,
                                         int *errno) {
   // Get ID for the vnode
   auto vnode_id = GetVnodeIDForVnode(ctx, vp);
-  if (!vnode_id) return KAUTH_RESULT_DEFER;
+  if (vnode_id.fsid == 0 && vnode_id.fileid == 0) return KAUTH_RESULT_DEFER;
 
   // Fetch decision
   auto returnedAction = FetchDecision(cred, vp, vnode_id);
 
-  // If file has dirty blocks, remove from cache and deny. This would usually
-  // be the case if a file has been written to and flushed but not yet
-  // closed.
-  if (vnode_hasdirtyblks(vp)) {
-    RemoveFromCache(vnode_id);
-    returnedAction = ACTION_RESPOND_DENY;
-
-    char path[MAXPATHLEN];
-    int len = MAXPATHLEN;
-    path[MAXPATHLEN - 1] = 0;
-    LOGW("file has dirty blocks: %s", vn_getpath(vp, path, &len) ? "unknown" : path);
-  }
-
   switch (returnedAction) {
     case ACTION_RESPOND_ALLOW: {
       auto proc = vfs_context_proc(ctx);

Source/santa-driver/SantaDecisionManager.h

@@ -64,7 +64,7 @@ class SantaDecisionManager : public OSObject {
   void ConnectClient(pid_t pid);
 
   /// Called by SantaDriverClient when a client disconnects
-  void DisconnectClient(bool itDied = false);
+  void DisconnectClient(bool itDied = false, pid_t pid = proc_selfpid());
 
   /// Returns whether a client is currently connected or not.
   bool ClientConnected() const;
@@ -86,7 +86,7 @@ class SantaDecisionManager : public OSObject {
   kern_return_t StopListener();
 
   /// Adds a decision to the cache, with a timestamp.
-  void AddToCache(uint64_t identifier,
+  void AddToCache(santa_vnode_id_t identifier,
                   const santa_action_t decision,
                   const uint64_t microsecs = GetCurrentUptime());
 
@@ -94,20 +94,30 @@ class SantaDecisionManager : public OSObject {
     Fetches a response from the cache, first checking to see if the entry 
     has expired.
   */
-  santa_action_t GetFromCache(uint64_t identifier);
+  santa_action_t GetFromCache(santa_vnode_id_t identifier);
 
   /// Checks to see if a given identifier is in the cache and removes it.
-  void RemoveFromCache(uint64_t identifier);
+  void RemoveFromCache(santa_vnode_id_t identifier);
 
   /// Returns the number of entries in the cache.
-  uint64_t RootCacheCount() const;
-  uint64_t NonRootCacheCount() const;
+  uint64_t CacheCount() const;
+
+  /// Clears the cache.
+  void ClearCache();
+
 
   /**
-    Clears the cache(s). If non_root_only is true, only the non-root cache 
-    is cleared.
+    Fills out the per_bucket_counts array with the number of items in each bucket in the
+    non-root decision cache.
+
+    @param per_bucket_counts An array of uint16_t's to fill in with the number of items in each
+        bucket. The size of this array is expected to equal array_size.
+    @param array_size The size of the per_bucket_counts array on input. Upon return this will be
+        updated to the number of slots that were actually used.
+    @param start_bucket If non-zero this is the bucket in the cache to start from. Upon return this
+        will be the next numbered bucket to start from for subsequent requests.
   */
-  void ClearCache(bool non_root_only = false);
+  void CacheBucketCount(uint16_t *per_bucket_counts, uint16_t *array_size, uint64_t *start_bucket);
 
   /// Increments the count of active callbacks pending.
   void IncrementListenerInvocations();
@@ -120,15 +130,18 @@ class SantaDecisionManager : public OSObject {
 
    @param ctx The VFS context to use.
    @param vp The Vnode to get the ID for
-   @return uint64_t The Vnode ID as a 64-bit unsigned int.
+   @return santa_vnode_id_t The Vnode ID.
    */
-  static inline uint64_t GetVnodeIDForVnode(const vfs_context_t ctx, const vnode_t vp) {
+  static inline santa_vnode_id_t GetVnodeIDForVnode(const vfs_context_t ctx, const vnode_t vp) {
     struct vnode_attr vap;
     VATTR_INIT(&vap);
     VATTR_WANTED(&vap, va_fsid);
     VATTR_WANTED(&vap, va_fileid);
     vnode_getattr(vp, &vap, ctx);
-    return (((uint64_t)vap.va_fsid << 32) | vap.va_fileid);
+    return {
+      .fsid = vap.va_fsid,
+      .fileid = vap.va_fileid
+    };
   }
 
   /**
@@ -198,7 +211,7 @@ class SantaDecisionManager : public OSObject {
     @param identifier The vnode ID string for this request
     @return santa_action_t The response for this request
   */
-  santa_action_t GetFromDaemon(santa_message_t *message, uint64_t identifier);
+  santa_action_t GetFromDaemon(santa_message_t *message, santa_vnode_id_t identifier);
 
   /**
     Fetches an execution decision for a file, first using the cache and then
@@ -212,7 +225,7 @@ class SantaDecisionManager : public OSObject {
     @return santa_action_t The response for this request
   */
   santa_action_t FetchDecision(
-      const kauth_cred_t cred, const vnode_t vp, const uint64_t vnode_id);
+      const kauth_cred_t cred, const vnode_t vp, const santa_vnode_id_t vnode_id);
 
   /**
     Posts the requested message to the decision data queue.
@@ -266,21 +279,8 @@ class SantaDecisionManager : public OSObject {
     return (uint64_t)((sec * 1000000) + usec);
   }
 
-  SantaCache<uint64_t> *root_decision_cache_;
-  SantaCache<uint64_t> *non_root_decision_cache_;
-  SantaCache<uint64_t> *vnode_pid_map_;
-
-  /**
-    Return the correct cache for a given identifier.
-
-    @param identifier The identifier
-    @return SantaCache* The cache to use
-  */
-  SantaCache<uint64_t>* CacheForIdentifier(const uint64_t identifier);
-
-  // This is the file system ID of the root filesystem,
-  // used to determine which cache to use for requests
-  uint32_t root_fsid_;
+  SantaCache<santa_vnode_id_t, uint64_t> *decision_cache_;
+  SantaCache<santa_vnode_id_t, uint64_t> *vnode_pid_map_;
 
   lck_grp_t *sdm_lock_grp_;
   lck_grp_attr_t *sdm_lock_grp_attr_;
@@ -301,7 +301,8 @@ class SantaDecisionManager : public OSObject {
   kauth_listener_t vnode_listener_;
   kauth_listener_t fileop_listener_;
 
-  struct timespec ts_;
+  struct timespec ts_= { .tv_sec = kRequestLoopSleepMilliseconds / 1000,
+                         .tv_nsec = kRequestLoopSleepMilliseconds % 1000 * 1000000 };
 };
 
 /**

Source/santa-driver/SantaDriverClient.cc

@@ -37,34 +37,38 @@ bool SantaDriverClient::initWithTask(
 
 bool SantaDriverClient::start(IOService *provider) {
   myProvider = OSDynamicCast(com_google_SantaDriver, provider);
-
   if (!myProvider) return false;
-  if (!super::start(provider)) return false;
 
   decisionManager = myProvider->GetDecisionManager();
   if (!decisionManager) return false;
+  decisionManager->retain();
 
-  return true;
+  return super::start(provider);
 }
 
 void SantaDriverClient::stop(IOService *provider) {
-  super::stop(provider);
   myProvider = nullptr;
+  decisionManager->release();
   decisionManager = nullptr;
+  super::stop(provider);
+}
+
+IOReturn SantaDriverClient::clientDied() {
+  LOGI("Client died.");
+  decisionManager->DisconnectClient(true);
+  return terminate(0) ? kIOReturnSuccess : kIOReturnError;
 }
 
 IOReturn SantaDriverClient::clientClose() {
-  decisionManager->DisconnectClient(true);
-  return terminate(kIOServiceSynchronous) ? kIOReturnSuccess : kIOReturnError;
+  LOGI("Client disconnected.");
+  decisionManager->DisconnectClient();
+  return terminate(0) ? kIOReturnSuccess : kIOReturnError;
 }
 
-bool SantaDriverClient::terminate(IOOptionBits options) {
-  decisionManager->DisconnectClient();
-  LOGI("Client disconnected.");
-
+bool SantaDriverClient::didTerminate(IOService *provider, IOOptionBits options, bool *defer) {
+  decisionManager->DisconnectClient(false, 0);
   if (myProvider && myProvider->isOpen(this)) myProvider->close(this);
-
-  return super::terminate(options);
+  return super::didTerminate(provider, options, defer);
 }
 
 #pragma mark Fetching memory and data queue notifications
@@ -133,9 +137,10 @@ IOReturn SantaDriverClient::allow_binary(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  const uint64_t vnode_id = static_cast<const uint64_t>(arguments->scalarInput[0]);
-  if (!vnode_id) return kIOReturnInvalid;
-  me->decisionManager->AddToCache(vnode_id, ACTION_RESPOND_ALLOW);
+  if (arguments->structureInputSize != sizeof(santa_vnode_id_t)) return kIOReturnInvalid;
+  santa_vnode_id_t *vnode_id = (santa_vnode_id_t *)arguments->structureInput;
+  if (vnode_id->fsid == 0 || vnode_id->fileid == 0) return kIOReturnInvalid;
+  me->decisionManager->AddToCache(*vnode_id, ACTION_RESPOND_ALLOW);
 
   return kIOReturnSuccess;
 }
@@ -145,9 +150,10 @@ IOReturn SantaDriverClient::deny_binary(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  const uint64_t vnode_id = static_cast<const uint64_t>(arguments->scalarInput[0]);
-  if (!vnode_id) return kIOReturnInvalid;
-  me->decisionManager->AddToCache(vnode_id, ACTION_RESPOND_DENY);
+  if (arguments->structureInputSize != sizeof(santa_vnode_id_t)) return kIOReturnInvalid;
+  santa_vnode_id_t *vnode_id = (santa_vnode_id_t *)arguments->structureInput;
+  if (vnode_id->fsid == 0 || vnode_id->fileid == 0) return kIOReturnInvalid;
+  me->decisionManager->AddToCache(*vnode_id, ACTION_RESPOND_DENY);
 
   return kIOReturnSuccess;
 }
@@ -157,9 +163,10 @@ IOReturn SantaDriverClient::acknowledge_binary(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  const uint64_t vnode_id = static_cast<const uint64_t>(arguments->scalarInput[0]);
-  if (!vnode_id) return kIOReturnInvalid;
-  me->decisionManager->AddToCache(vnode_id, ACTION_RESPOND_ACK, 0);
+  if (arguments->structureInputSize != sizeof(santa_vnode_id_t)) return kIOReturnInvalid;
+  santa_vnode_id_t *vnode_id = (santa_vnode_id_t *)arguments->structureInput;
+  if (vnode_id->fsid == 0 || vnode_id->fileid == 0) return kIOReturnInvalid;
+  me->decisionManager->AddToCache(*vnode_id, ACTION_RESPOND_ACK);
 
   return kIOReturnSuccess;
 }
@@ -169,8 +176,7 @@ IOReturn SantaDriverClient::clear_cache(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  const bool non_root_only = static_cast<const bool>(arguments->scalarInput[0]);
-  me->decisionManager->ClearCache(non_root_only);
+  me->decisionManager->ClearCache();
 
   return kIOReturnSuccess;
 }
@@ -180,8 +186,7 @@ IOReturn SantaDriverClient::cache_count(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  arguments->scalarOutput[0] = me->decisionManager->RootCacheCount();
-  arguments->scalarOutput[1] = me->decisionManager->NonRootCacheCount();
+  arguments->scalarOutput[0] = me->decisionManager->CacheCount();
   return kIOReturnSuccess;
 }
 
@@ -190,8 +195,27 @@ IOReturn SantaDriverClient::check_cache(
   SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
   if (!me) return kIOReturnBadArgument;
 
-  const uint64_t input = static_cast<const uint64_t>(arguments->scalarInput[0]);
-  arguments->scalarOutput[0] = me->decisionManager->GetFromCache(input);
+  if (arguments->structureInputSize != sizeof(santa_vnode_id_t)) return kIOReturnInvalid;
+  santa_vnode_id_t *vnode_id = (santa_vnode_id_t *)arguments->structureInput;
+  if (vnode_id->fsid == 0 || vnode_id->fileid == 0) return kIOReturnInvalid;
+  arguments->scalarOutput[0] = me->decisionManager->GetFromCache(*vnode_id);
+
+  return kIOReturnSuccess;
+}
+
+IOReturn SantaDriverClient::cache_bucket_count(
+    OSObject *target, void *reference, IOExternalMethodArguments *arguments) {
+  SantaDriverClient *me = OSDynamicCast(SantaDriverClient, target);
+  if (!me) return kIOReturnBadArgument;
+
+  santa_bucket_count_t *counts = reinterpret_cast<santa_bucket_count_t *>(
+      arguments->structureOutput);
+  const santa_bucket_count_t *input = reinterpret_cast<const santa_bucket_count_t *>(
+      arguments->structureInput);
+
+  uint16_t s = sizeof(counts->per_bucket) / sizeof(uint16_t);
+  counts->start = input->start;
+  me->decisionManager->CacheBucketCount(counts->per_bucket, &s, &(counts->start));
 
   return kIOReturnSuccess;
 }
@@ -209,12 +233,14 @@ IOReturn SantaDriverClient::externalMethod(
   static IOExternalMethodDispatch sMethods[kSantaUserClientNMethods] = {
     // Function ptr, input scalar count, input struct size, output scalar count, output struct size
     { &SantaDriverClient::open, 0, 0, 0, 0 },
-    { &SantaDriverClient::allow_binary, 1, 0, 0, 0 },
-    { &SantaDriverClient::deny_binary, 1, 0, 0, 0 },
-    { &SantaDriverClient::acknowledge_binary, 1, 0, 0, 0 },
-    { &SantaDriverClient::clear_cache, 1, 0, 0, 0 },
-    { &SantaDriverClient::cache_count, 0, 0, 2, 0 },
-    { &SantaDriverClient::check_cache, 1, 0, 1, 0 }
+    { &SantaDriverClient::allow_binary, 0, sizeof(santa_vnode_id_t), 0, 0 },
+    { &SantaDriverClient::deny_binary, 0, sizeof(santa_vnode_id_t), 0, 0 },
+    { &SantaDriverClient::acknowledge_binary, 0, sizeof(santa_vnode_id_t), 0, 0 },
+    { &SantaDriverClient::clear_cache, 0, 0, 0, 0 },
+    { &SantaDriverClient::cache_count, 0, 0, 1, 0 },
+    { &SantaDriverClient::check_cache, 0, sizeof(santa_vnode_id_t), 1, 0 },
+    { &SantaDriverClient::cache_bucket_count, 0, sizeof(santa_bucket_count_t),
+        0, sizeof(santa_bucket_count_t) },
   };
 
   if (selector > static_cast<UInt32>(kSantaUserClientNMethods)) {

Source/santa-driver/SantaDriverClient.h

@@ -47,11 +47,14 @@ class com_google_SantaDriverClient : public IOUserClient {
   ///  Called when this class is stopping
   void stop(IOService *provider) override;
 
-  ///  Called when a client disconnects
+  ///  Called when a client manually disconnects (via IOServiceClose)
   IOReturn clientClose() override;
 
-  ///  Called when the driver is shutting down
-  bool terminate(IOOptionBits options) override;
+  ///  Called when a client dies
+  IOReturn clientDied() override;
+
+  ///  Called during termination
+  bool didTerminate(IOService* provider, IOOptionBits options, bool* defer) override;
 
   ///  Called in clients with IOConnectSetNotificationPort
   IOReturn registerNotificationPort(
@@ -105,6 +108,11 @@ class com_google_SantaDriverClient : public IOUserClient {
   static IOReturn check_cache(
       OSObject *target, void *reference, IOExternalMethodArguments *arguments);
 
+  ///  The daemon calls this to find out how many items are in each cache bucket.
+  ///  Input and output are both an instance of santa_bucket_count_t.
+  static IOReturn cache_bucket_count(
+      OSObject *target, void *reference, IOExternalMethodArguments *arguments);
+
  private:
   com_google_SantaDriver *myProvider;
   SantaDecisionManager *decisionManager;

Source/santactl/Commands/SNTCommandBundleInfo.m

@@ -22,14 +22,14 @@
 #import "SNTStoredEvent.h"
 #import "SNTXPCControlInterface.h"
 
+#ifdef DEBUG
+
 @interface SNTCommandBundleInfo : SNTCommand<SNTCommandProtocol>
 @end
 
 @implementation SNTCommandBundleInfo
 
-#ifdef DEBUG
 REGISTER_COMMAND_NAME(@"bundleinfo")
-#endif
 
 + (BOOL)requiresRoot {
   return NO;
@@ -78,3 +78,5 @@ REGISTER_COMMAND_NAME(@"bundleinfo")
 }
 
 @end
+
+#endif

Source/santactl/Commands/SNTCommandCacheHistogram.m

@@ -0,0 +1,81 @@
+/// Copyright 2018 Google Inc. All rights reserved.
+///
+/// Licensed under the Apache License, Version 2.0 (the "License");
+/// you may not use this file except in compliance with the License.
+/// You may obtain a copy of the License at
+///
+///    http://www.apache.org/licenses/LICENSE-2.0
+///
+///    Unless required by applicable law or agreed to in writing, software
+///    distributed under the License is distributed on an "AS IS" BASIS,
+///    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+///    See the License for the specific language governing permissions and
+///    limitations under the License.
+
+#ifdef DEBUG
+
+#import <Foundation/Foundation.h>
+
+#import "SNTCommand.h"
+#import "SNTCommandController.h"
+
+#import <MOLXPCConnection/MOLXPCConnection.h>
+
+#import "SNTLogging.h"
+#import "SNTXPCControlInterface.h"
+
+@interface SNTCommandCacheHistogram : SNTCommand<SNTCommandProtocol>
+@end
+
+@implementation SNTCommandCacheHistogram
+
+REGISTER_COMMAND_NAME(@"cachehistogram")
+
++ (BOOL)requiresRoot {
+  return YES;
+}
+
++ (BOOL)requiresDaemonConn {
+  return YES;
+}
+
++ (NSString *)shortHelpText {
+  return @"Print a cache distribution histogram.";
+}
+
++ (NSString *)longHelpText {
+  return (@"Prints a histogram of each bucket of the in-kernel cache\n"
+          @"  Use -g to get 'graphical' output\n"
+          @"Only available in DEBUG builds.");
+}
+
+- (void)runWithArguments:(NSArray *)arguments {
+  [[self.daemonConn remoteObjectProxy] cacheBucketCount:^(NSArray *counts) {
+    NSMutableDictionary<NSNumber *, NSNumber *> *d = [NSMutableDictionary dictionary];
+    [counts enumerateObjectsUsingBlock:^(id  _Nonnull obj, NSUInteger idx, BOOL * _Nonnull stop) {
+      d[obj] = @([d[obj] intValue] + 1);
+    }];
+    printf("There are %llu empty buckets\n", [d[@0] unsignedLongLongValue]);
+
+    for (NSNumber *key in [d.allKeys sortedArrayUsingSelector:@selector(compare:)]) {
+      if ([key isEqual:@0]) continue;
+      uint64_t k = [key unsignedLongLongValue];
+      uint64_t v = [d[key] unsignedLongLongValue];
+
+      if ([[[NSProcessInfo processInfo] arguments] containsObject:@"-g"]) {
+        printf("%4llu: ", k);
+        for (uint64_t y = 0; y < v; ++y) {
+          printf("#");
+        }
+        printf("\n");
+      } else {
+        printf("%4llu: %llu\n", k, v);
+      }
+    }
+    exit(0);
+  }];
+}
+
+@end
+
+#endif

Source/santactl/Commands/SNTCommandCheckCache.m

@@ -24,14 +24,14 @@
 
 #include <sys/stat.h>
 
+#ifdef DEBUG
+
 @interface SNTCommandCheckCache : SNTCommand<SNTCommandProtocol>
 @end
 
 @implementation SNTCommandCheckCache
 
-#ifdef DEBUG
 REGISTER_COMMAND_NAME(@"checkcache")
-#endif
 
 + (BOOL)requiresRoot {
   return NO;
@@ -51,7 +51,7 @@ REGISTER_COMMAND_NAME(@"checkcache")
 }
 
 - (void)runWithArguments:(NSArray *)arguments {
-  uint64_t vnodeID = [self vnodeIDForFile:arguments.firstObject];
+  santa_vnode_id_t vnodeID = [self vnodeIDForFile:arguments.firstObject];
   [[self.daemonConn remoteObjectProxy] checkCacheForVnodeID:vnodeID
                                                   withReply:^(santa_action_t action) {
     if (action == ACTION_RESPOND_ALLOW) {
@@ -67,10 +67,13 @@ REGISTER_COMMAND_NAME(@"checkcache")
   }];
 }
 
-- (uint64_t)vnodeIDForFile:(NSString *)path {
+- (santa_vnode_id_t)vnodeIDForFile:(NSString *)path {
   struct stat fstat = {};
   stat(path.fileSystemRepresentation, &fstat);
-  return (((uint64_t)fstat.st_dev << 32) | fstat.st_ino);
+  santa_vnode_id_t ret = {.fsid = fstat.st_dev, .fileid = fstat.st_ino};
+  return ret;
 }
 
 @end
+
+#endif

Source/santactl/Commands/SNTCommandFlushCache.m

@@ -12,6 +12,8 @@
 ///    See the License for the specific language governing permissions and
 ///    limitations under the License.
 
+#ifdef DEBUG
+
 #import <Foundation/Foundation.h>
 
 #import "SNTCommand.h"
@@ -27,9 +29,7 @@
 
 @implementation SNTCommandFlushCache
 
-#ifdef DEBUG
 REGISTER_COMMAND_NAME(@"flushcache")
-#endif
 
 + (BOOL)requiresRoot {
   return YES;
@@ -61,3 +61,5 @@ REGISTER_COMMAND_NAME(@"flushcache")
 }
 
 @end
+
+#endif

Source/santactl/Commands/SNTCommandStatus.m

@@ -50,10 +50,16 @@ REGISTER_COMMAND_NAME(@"status")
   dispatch_group_t group = dispatch_group_create();
 
   // Daemon status
+  __block BOOL driverConnected;
   __block NSString *clientMode;
   __block uint64_t cpuEvents, ramEvents;
   __block double cpuPeak, ramPeak;
   dispatch_group_enter(group);
+  [[self.daemonConn remoteObjectProxy] driverConnectionEstablished:^(BOOL connected) {
+    driverConnected = connected;
+    dispatch_group_leave(group);
+  }];
+  dispatch_group_enter(group);
   [[self.daemonConn remoteObjectProxy] clientMode:^(SNTClientMode cm) {
     switch (cm) {
       case SNTClientModeMonitor:
@@ -81,11 +87,10 @@ REGISTER_COMMAND_NAME(@"status")
   BOOL fileLogging = ([[SNTConfigurator configurator] fileChangesRegex] != nil);
 
   // Kext status
-  __block uint64_t rootCacheCount = -1, nonRootCacheCount = -1;
+  __block uint64_t cacheCount = -1;
   dispatch_group_enter(group);
-  [[self.daemonConn remoteObjectProxy] cacheCounts:^(uint64_t rootCache, uint64_t nonRootCache) {
-    rootCacheCount = rootCache;
-    nonRootCacheCount = nonRootCache;
+  [[self.daemonConn remoteObjectProxy] cacheCounts:^(uint64_t count) {
+    cacheCount = count;
     dispatch_group_leave(group);
   }];
 
@@ -160,6 +165,7 @@ REGISTER_COMMAND_NAME(@"status")
   if ([arguments containsObject:@"--json"]) {
     NSDictionary *stats = @{
       @"daemon" : @{
+        @"driver_connected" : @(driverConnected),
         @"mode" : clientMode ?: @"null",
         @"file_logging" : @(fileLogging),
         @"watchdog_cpu_events" : @(cpuEvents),
@@ -168,8 +174,7 @@ REGISTER_COMMAND_NAME(@"status")
         @"watchdog_ram_peak" : @(ramPeak),
       },
       @"kernel" : @{
-        @"root_cache_count" : @(rootCacheCount),
-        @"non_root_cache_count": @(nonRootCacheCount),
+        @"cache_count" : @(cacheCount),
       },
       @"database" : @{
         @"binary_rules" : @(binaryRuleCount),
@@ -192,13 +197,13 @@ REGISTER_COMMAND_NAME(@"status")
     printf("%s\n", [statsStr UTF8String]);
   } else {
     printf(">>> Daemon Info\n");
+    printf("  %-25s | %s\n", "Driver Connected", driverConnected ? "Yes" : "No");
     printf("  %-25s | %s\n", "Mode", [clientMode UTF8String]);
     printf("  %-25s | %s\n", "File Logging", (fileLogging ? "Yes" : "No"));
     printf("  %-25s | %lld  (Peak: %.2f%%)\n", "Watchdog CPU Events", cpuEvents, cpuPeak);
     printf("  %-25s | %lld  (Peak: %.2fMB)\n", "Watchdog RAM Events", ramEvents, ramPeak);
     printf(">>> Kernel Info\n");
-    printf("  %-25s | %lld\n", "Root cache count", rootCacheCount);
-    printf("  %-25s | %lld\n", "Non-root cache count", nonRootCacheCount);
+    printf("  %-25s | %lld\n", "Cache count", cacheCount);
     printf(">>> Database Info\n");
     printf("  %-25s | %lld\n", "Binary Rules", binaryRuleCount);
     printf("  %-25s | %lld\n", "Certificate Rules", certRuleCount);

Source/santad/Logs/SNTEventLog.h

@@ -50,4 +50,5 @@
 
 // A UTC Date formatter.
 @property(readonly, nonatomic) NSDateFormatter *dateFormatter;
+@property(readonly, nonatomic) NSString *machineID;
 @end

Source/santad/Logs/SNTEventLog.m

@@ -20,6 +20,7 @@
 #include <sys/sysctl.h>
 
 #import "SNTCachedDecision.h"
+#import "SNTConfigurator.h"
 
 @interface SNTEventLog ()
 @property NSMutableDictionary<NSNumber *, SNTCachedDecision *> *detailStore;
@@ -43,6 +44,9 @@
     _dateFormatter = [[NSDateFormatter alloc] init];
     _dateFormatter.dateFormat = @"yyyy-MM-dd'T'HH:mm:ss.SSS'Z'";
     _dateFormatter.timeZone = [NSTimeZone timeZoneWithName:@"UTC"];
+
+    // Grab the system UUID on init
+     _machineID = [[SNTConfigurator configurator] machineID];
   }
   return self;
 }
@@ -77,14 +81,14 @@
 
 - (void)cacheDecision:(SNTCachedDecision *)cd {
   dispatch_sync(self.detailStoreQueue, ^{
-    self.detailStore[@(cd.vnodeId)] = cd;
+    self.detailStore[@(cd.vnodeId.fileid)] = cd;
   });
 }
 
 - (SNTCachedDecision *)cachedDecisionForMessage:(santa_message_t)message {
   __block SNTCachedDecision *cd;
   dispatch_sync(self.detailStoreQueue, ^{
-    cd = self.detailStore[@(message.vnode_id)];
+    cd = self.detailStore[@(message.vnode_id.fileid)];
   });
   return cd;
 }

Source/santad/Logs/SNTSyslogEventLog.m

@@ -69,6 +69,10 @@
       message.uid, [self nameForUID:message.uid],
       message.gid, [self nameForGID:message.gid]];
 
+  if ([[SNTConfigurator configurator] enableMachineIDDecoration]) {
+    [outStr appendFormat:@"|machineid=%@", self.machineID];
+  }
+
   [self writeLog:outStr];
 }
 
@@ -165,6 +169,10 @@
     [self addArgsForPid:message.pid toString:outLog];
   }
 
+  if ([[SNTConfigurator configurator] enableMachineIDDecoration]) {
+    [outLog appendFormat:@"|machineid=%@", self.machineID];
+  }
+
   [self writeLog:outLog];
 }
 

Source/santad/SNTApplication.m

@@ -250,7 +250,7 @@ void diskDisappearedCallback(DADiskRef disk, void *context) {
   if (![props[@"DAVolumeMountable"] boolValue]) return;
 
   [app.eventLog logDiskDisappeared:props];
-  [app.driverManager flushCacheNonRootOnly:YES];
+  [app.driverManager flushCache];
 }
 
 - (void)startSyncd {
@@ -303,7 +303,7 @@ void diskDisappearedCallback(DADiskRef disk, void *context) {
     if (!new && !old) return;
     if (![new.pattern isEqualToString:old.pattern]) {
       LOGI(@"Changed [white|black]list regex, flushing cache");
-      [self.driverManager flushCacheNonRootOnly:NO];
+      [self.driverManager flushCache];
     }
   }
 }
@@ -311,7 +311,7 @@ void diskDisappearedCallback(DADiskRef disk, void *context) {
 - (void)clientModeDidChange:(SNTClientMode)clientMode {
   if (clientMode == SNTClientModeLockdown) {
     LOGI(@"Changed client mode, flushing cache.");
-    [self.driverManager flushCacheNonRootOnly:NO];
+    [self.driverManager flushCache];
   }
   [[self.notQueue.notifierConnection remoteObjectProxy] postClientModeNotification:clientMode];
 }

Source/santad/SNTCachedDecision.h

@@ -15,13 +15,14 @@
 #import <Foundation/Foundation.h>
 
 #import "SNTCommonEnums.h"
+#import "SNTKernelCommon.h"
 
 ///
 ///  Store information about executions from decision making for later logging.
 ///
 @interface SNTCachedDecision : NSObject
 
-@property uint64_t vnodeId;
+@property santa_vnode_id_t vnodeId;
 @property SNTEventState decision;
 @property NSString *decisionExtra;
 @property NSString *sha256;

Source/santad/SNTDaemonControlController.m

@@ -70,19 +70,27 @@ double watchdogRAMPeak = 0;
 
 #pragma mark Kernel ops
 
-- (void)cacheCounts:(void (^)(uint64_t, uint64_t))reply {
-  NSArray<NSNumber *> *counts = [self.driverManager cacheCounts];
-  reply([counts[0] unsignedLongLongValue], [counts[1] unsignedLongLongValue]);
+- (void)cacheCounts:(void (^)(uint64_t))reply {
+  uint64_t count = [self.driverManager cacheCount];
+  reply(count);
+}
+
+- (void)cacheBucketCount:(void (^)(NSArray *))reply {
+  reply([self.driverManager cacheBucketCount]);
 }
 
 - (void)flushCache:(void (^)(BOOL))reply {
-  reply([self.driverManager flushCacheNonRootOnly:NO]);
+  reply([self.driverManager flushCache]);
 }
 
-- (void)checkCacheForVnodeID:(uint64_t)vnodeID withReply:(void (^)(santa_action_t))reply {
+- (void)checkCacheForVnodeID:(santa_vnode_id_t)vnodeID withReply:(void (^)(santa_action_t))reply {
   reply([self.driverManager checkCache:vnodeID]);
 }
 
+- (void)driverConnectionEstablished:(void (^)(BOOL))reply {
+  reply(self.driverManager.connectionEstablished);
+}
+
 #pragma mark Database ops
 
 - (void)databaseRuleCounts:(void (^)(int64_t binary, int64_t certificate))reply {
@@ -100,7 +108,7 @@ double watchdogRAMPeak = 0;
   NSPredicate *p = [NSPredicate predicateWithFormat:@"SELF.state != %d", SNTRuleStateWhitelist];
   if ([rules filteredArrayUsingPredicate:p].count || cleanSlate) {
     LOGI(@"Received non-whitelist rule, flushing cache");
-    [self.driverManager flushCacheNonRootOnly:NO];
+    [self.driverManager flushCache];
   }
 
   reply(error);

Source/santad/SNTDriverManager.h

@@ -45,21 +45,30 @@
 ///
 ///  Sends a response to a query back to the kernel.
 ///
-- (kern_return_t)postToKernelAction:(santa_action_t)action forVnodeID:(uint64_t)vnodeId;
+- (kern_return_t)postToKernelAction:(santa_action_t)action forVnodeID:(santa_vnode_id_t)vnodeId;
 
 ///
 ///  Get the number of binaries in the kernel's caches.
 ///
-- (NSArray<NSNumber *> *)cacheCounts;
+- (uint64_t)cacheCount;
+
+///
+///  Return an array representing all buckets in the kernel's decision cache where each number
+///  is the number of items in that bucket.
+///
+- (NSArray<NSNumber *> *)cacheBucketCount;
 
 ///
 ///  Flush the kernel's binary cache.
 ///
-- (BOOL)flushCacheNonRootOnly:(BOOL)nonRootOnly;
+- (BOOL)flushCache;
 
 ///
 ///  Check the kernel cache for a VnodeID
 ///
-- (santa_action_t)checkCache:(uint64_t)vnodeID;
+- (santa_action_t)checkCache:(santa_vnode_id_t)vnodeID;
+
+///  Returns whether the connection to the driver has been established.
+@property(readonly) BOOL connectionEstablished;
 
 @end

Source/santad/SNTDriverManager.m

@@ -23,61 +23,27 @@
 
 @interface SNTDriverManager ()
 @property io_connect_t connection;
+@property(readwrite) BOOL connectionEstablished;
 @end
 
 @implementation SNTDriverManager
 
-static const int MAX_DELAY = 15;
-
 #pragma mark init/dealloc
 
 - (instancetype)init {
   self = [super init];
   if (self) {
-    kern_return_t kr;
-    io_service_t serviceObject;
-    CFDictionaryRef classToMatch;
-
-    if (!(classToMatch = IOServiceMatching(USERCLIENT_CLASS))) {
-      LOGD(@"Failed to create matching dictionary");
+    CFDictionaryRef classToMatch = IOServiceMatching(USERCLIENT_CLASS);
+    if (!classToMatch) {
+      LOGE(@"Failed to create matching dictionary");
       return nil;
     }
 
     // Attempt to load driver. It may already be running, so ignore any return value.
     KextManagerLoadKextWithIdentifier(CFSTR(USERCLIENT_ID), NULL);
 
-    // Locate driver. Wait for it if necessary.
-    int delay = 1;
-    do {
-      CFRetain(classToMatch);  // this ref is released by IOServiceGetMatchingService
-      serviceObject = IOServiceGetMatchingService(kIOMasterPortDefault, classToMatch);
-
-      if (!serviceObject) {
-        LOGD(@"Waiting for Santa driver to become available");
-        sleep(delay);
-        if (delay < MAX_DELAY) delay *= 2;
-      }
-    } while (!serviceObject);
-    CFRelease(classToMatch);
-
-    // This calls `initWithTask`, `attach` and `start` in `SantaDriverClient`
-    kr = IOServiceOpen(serviceObject, mach_task_self(), 0, &_connection);
-    IOObjectRelease(serviceObject);
-    if (kr != kIOReturnSuccess) {
-      LOGD(@"Failed to open Santa driver service");
-      return nil;
-    }
-
-    // Call `open` in `SantaDriverClient`
-    kr = IOConnectCallMethod(_connection, kSantaUserClientOpen, 0, 0, 0, 0, 0, 0, 0, 0);
-
-    if (kr == kIOReturnExclusiveAccess) {
-      LOGD(@"A client is already connected");
-      return nil;
-    } else if (kr != kIOReturnSuccess) {
-      LOGD(@"An error occurred while opening the connection");
-      return nil;
-    }
+    // Wait for the driver to appear
+    [self waitForDriver:classToMatch];
   }
   return self;
 }
@@ -86,6 +52,71 @@ static const int MAX_DELAY = 15;
   IOServiceClose(_connection);
 }
 
+#pragma mark Driver Waiting
+
+// Helper function used with IOServiceAddMatchingNotification which expects
+// a DriverAppearedBlock to be passed as the 'reference' argument. The block
+// will be called for each object iterated over and if the block wants to keep
+// the object, it should call IOObjectRetain().
+typedef void (^DriverAppearedBlock)(io_object_t object);
+static void driverAppearedHandler(void *info, io_iterator_t iterator) {
+  DriverAppearedBlock block = (__bridge DriverAppearedBlock)info;
+  if (!block) return;
+  io_object_t object = 0;
+  while ((object = IOIteratorNext(iterator))) {
+    block(object);
+    IOObjectRelease(object);
+  }
+}
+
+// Wait for the driver to appear, then attach to and open it.
+- (void)waitForDriver:(CFDictionaryRef CF_RELEASES_ARGUMENT)matchingDict {
+  IONotificationPortRef notificationPort = IONotificationPortCreate(kIOMasterPortDefault);
+  CFRunLoopAddSource([[NSRunLoop currentRunLoop] getCFRunLoop],
+                     IONotificationPortGetRunLoopSource(notificationPort),
+                     kCFRunLoopDefaultMode);
+
+  io_iterator_t iterator = 0;
+
+  DriverAppearedBlock block = ^(io_object_t object) {
+    // This calls `initWithTask`, `attach` and `start` in `SantaDriverClient`
+    kern_return_t kr = IOServiceOpen(object, mach_task_self(), 0, &_connection);
+    if (kr != kIOReturnSuccess) {
+      LOGE(@"Failed to open santa-driver service: 0x%X", kr);
+      exit(1);
+    }
+
+    // Call `open` in `SantaDriverClient`
+    kr = IOConnectCallMethod(_connection, kSantaUserClientOpen, 0, 0, 0, 0, 0, 0, 0, 0);
+    if (kr == kIOReturnExclusiveAccess) {
+      LOGE(@"A client is already connected");
+      exit(2);
+    } else if (kr != kIOReturnSuccess) {
+      LOGE(@"An error occurred while opening the connection: 0x%X", kr);
+      exit(3);
+    }
+
+    // Release the iterator to disarm the notifications.
+    IOObjectRelease(iterator);
+    IONotificationPortDestroy(notificationPort);
+
+    _connectionEstablished = YES;
+  };
+
+  LOGI(@"Waiting for Santa driver to become available");
+
+  IOServiceAddMatchingNotification(notificationPort,
+                                   kIOMatchedNotification,
+                                   matchingDict,
+                                   driverAppearedHandler,
+                                   (__bridge_retained void *)block,
+                                   &iterator);
+
+  // Call the handler once to 'empty' the iterator, arming it. If the driver is already loaded
+  // this will immediately cause the connection to be fully established.
+  driverAppearedHandler((__bridge_retained void*)block, iterator);
+}
+
 #pragma mark Incoming messages
 
 - (void)listenForDecisionRequests:(void (^)(santa_message_t))callback {
@@ -98,7 +129,7 @@ static const int MAX_DELAY = 15;
 
 - (void)listenForRequestsOfType:(santa_queuetype_t)type
                    withCallback:(void (^)(santa_message_t))callback {
-  kern_return_t kr;
+  while (!self.connectionEstablished) usleep(100000); // 100ms
 
   // Allocate a mach port to receive notifactions from the IODataQueue
   mach_port_t receivePort = IODataQueueAllocateNotificationPort();
@@ -108,9 +139,9 @@ static const int MAX_DELAY = 15;
   }
 
   // This will call registerNotificationPort() inside our user client class
-  kr = IOConnectSetNotificationPort(self.connection, type, receivePort, 0);
+  kern_return_t kr = IOConnectSetNotificationPort(self.connection, type, receivePort, 0);
   if (kr != kIOReturnSuccess) {
-    LOGD(@"Failed to register notification port for type %d: %d", type, kr);
+    LOGD(@"Failed to register notification port for type %d: 0x%X", type, kr);
     mach_port_destroy(mach_task_self(), receivePort);
     return;
   }
@@ -120,7 +151,7 @@ static const int MAX_DELAY = 15;
   mach_vm_size_t size = 0;
   kr = IOConnectMapMemory(self.connection, type, mach_task_self(), &address, &size, kIOMapAnywhere);
   if (kr != kIOReturnSuccess) {
-    LOGD(@"Failed to map memory for type %d: %d", type, kr);
+    LOGD(@"Failed to map memory for type %d: 0x%X", type, kr);
     mach_port_destroy(mach_task_self(), receivePort);
     return;
   }
@@ -135,7 +166,7 @@ static const int MAX_DELAY = 15;
       if (kr == kIOReturnSuccess) {
         callback(vdata);
       } else {
-        LOGE(@"Error dequeuing data for type %d: %d", type, kr);
+        LOGE(@"Error dequeuing data for type %d: 0x%X", type, kr);
         exit(2);
       }
     }
@@ -147,27 +178,27 @@ static const int MAX_DELAY = 15;
 
 #pragma mark Outgoing messages
 
-- (kern_return_t)postToKernelAction:(santa_action_t)action forVnodeID:(uint64_t)vnodeId {
+- (kern_return_t)postToKernelAction:(santa_action_t)action forVnodeID:(santa_vnode_id_t)vnodeId {
   switch (action) {
     case ACTION_RESPOND_ALLOW:
-      return IOConnectCallScalarMethod(_connection,
+      return IOConnectCallStructMethod(_connection,
                                        kSantaUserClientAllowBinary,
                                        &vnodeId,
-                                       1,
+                                       sizeof(vnodeId),
                                        0,
                                        0);
     case ACTION_RESPOND_DENY:
-      return IOConnectCallScalarMethod(_connection,
+      return IOConnectCallStructMethod(_connection,
                                        kSantaUserClientDenyBinary,
                                        &vnodeId,
-                                       1,
+                                       sizeof(vnodeId),
                                        0,
                                        0);
     case ACTION_RESPOND_ACK:
-      return IOConnectCallScalarMethod(_connection,
+      return IOConnectCallStructMethod(_connection,
                                        kSantaUserClientAcknowledgeBinary,
                                        &vnodeId,
-                                       1,
+                                       sizeof(vnodeId),
                                        0,
                                        0);
     default:
@@ -175,9 +206,9 @@ static const int MAX_DELAY = 15;
   }
 }
 
-- (NSArray<NSNumber *> *)cacheCounts {
-  uint32_t input_count = 2;
-  uint64_t cache_counts[2] = {0, 0};
+- (uint64_t)cacheCount {
+  uint32_t input_count = 1;
+  uint64_t cache_counts[1] = {0};
 
   IOConnectCallScalarMethod(_connection,
                             kSantaUserClientCacheCount,
@@ -186,30 +217,47 @@ static const int MAX_DELAY = 15;
                             cache_counts,
                             &input_count);
 
-  return @[ @(cache_counts[0]), @(cache_counts[1]) ];
+  return cache_counts[0];
 }
 
-- (BOOL)flushCacheNonRootOnly:(BOOL)nonRootOnly {
-  const uint64_t nonRoot = nonRootOnly;
+- (BOOL)flushCache {
   return IOConnectCallScalarMethod(_connection,
                                    kSantaUserClientClearCache,
-                                   &nonRoot,
-                                   1,
+                                   0,
+                                   0,
                                    0,
                                    0) == KERN_SUCCESS;
 }
 
-- (santa_action_t)checkCache:(uint64_t)vnodeID {
-  uint32_t input_count = 1;
-  uint64_t vnode_action = 0;
+- (santa_action_t)checkCache:(santa_vnode_id_t)vnodeID {
+  uint64_t output;
+  uint32_t outputCnt = 1;
 
-  IOConnectCallScalarMethod(_connection,
-                            kSantaUserClientCheckCache,
-                            &vnodeID,
-                            1,
-                            &vnode_action,
-                            &input_count);
-  return (santa_action_t)vnode_action;
+  IOConnectCallMethod(self.connection, kSantaUserClientCheckCache,
+                      NULL, 0, &vnodeID, sizeof(santa_vnode_id_t),
+                      &output, &outputCnt, NULL, 0);
+  return (santa_action_t)output;
+}
+
+- (NSArray *)cacheBucketCount {
+  santa_bucket_count_t counts = {};
+  size_t size = sizeof(counts);
+
+  NSMutableArray *a = [NSMutableArray array];
+
+  do {
+    IOConnectCallStructMethod(self.connection,
+                              kSantaUserClientCacheBucketCount,
+                              &counts,
+                              size,
+                              &counts,
+                              &size);
+    for (uint64_t i = 0; i < sizeof(counts.per_bucket) / sizeof(uint16_t); ++i) {
+      [a addObject:@(counts.per_bucket[i])];
+    }
+  } while (counts.start > 0);
+
+  return a;
 }
 
 @end

Source/santad/SNTExecutionController.m

@@ -90,7 +90,7 @@ static size_t kLargeBinarySize = 30 * 1024 * 1024;
 - (void)validateBinaryWithMessage:(santa_message_t)message {
   // Get info about the file. If we can't get this info, allow execution and log an error.
   if (unlikely(message.path == NULL)) {
-    LOGE(@"Path for vnode_id is NULL: %llu", message.vnode_id);
+    LOGE(@"Path for vnode_id is NULL: %llu/%llu", message.vnode_id.fsid, message.vnode_id.fileid);
     [_driverManager postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:message.vnode_id];
     return;
   }
@@ -115,14 +115,16 @@ static size_t kLargeBinarySize = 30 * 1024 * 1024;
     [_driverManager postToKernelAction:ACTION_RESPOND_ACK forVnodeID:message.vnode_id];
   }
 
-  // Get codesigning info about the file.
-  NSError *csError;
-  MOLCodesignChecker *csInfo =
-      [[MOLCodesignChecker alloc] initWithBinaryPath:binInfo.path
-                                      fileDescriptor:binInfo.fileHandle.fileDescriptor
-                                               error:&csError];
-  // Ignore codesigning if there are any errors with the signature.
-  if (csError) csInfo = nil;
+  // Get codesigning info about the file but only if it's a Mach-O.
+  MOLCodesignChecker *csInfo;
+  if (binInfo.isMachO) {
+    NSError *csError;
+    csInfo = [[MOLCodesignChecker alloc] initWithBinaryPath:binInfo.path
+                                             fileDescriptor:binInfo.fileHandle.fileDescriptor
+                                                      error:&csError];
+    // Ignore codesigning if there are any errors with the signature.
+    if (csError) csInfo = nil;
+  }
 
   // Actually make the decision.
   SNTCachedDecision *cd = [self.policyProcessor decisionForFileInfo:binInfo
@@ -139,7 +141,7 @@ static size_t kLargeBinarySize = 30 * 1024 * 1024;
   if (action == ACTION_RESPOND_ALLOW) [_eventLog cacheDecision:cd];
 
   // Send the decision to the kernel.
-  [_driverManager postToKernelAction:action forVnodeID:cd.vnodeId];
+  [_driverManager postToKernelAction:action forVnodeID:message.vnode_id];
 
   // Log to database if necessary.
   if (cd.decision != SNTEventStateAllowBinary &&

Tests/KernelTests/main.mm

@@ -14,12 +14,14 @@
 
 #import <Foundation/Foundation.h>
 #import <IOKit/IOKitLib.h>
+#import <IOKit/kext/KextManager.h>
 
 #import <CommonCrypto/CommonDigest.h>
 
 #include <cmath>
 #include <ctime>
 #include <iostream>
+#include <libkern/OSKextLib.h>
 #include <mach/mach.h>
 #include <numeric>
 #include <sys/ptrace.h>
@@ -31,34 +33,41 @@
 ///
 /// Kernel Extension Tests
 ///
-/// Build and launch as root while the kernel extension is loaded and nothing is already connected.
+/// Build and launch as root. This target is dependent on the santa-driver target and these
+/// tests will load santa-driver from the same location this binary is executed from, unloading
+/// any existing driver (and daemon) if necessary.
 ///
 
 #define TSTART(testName) \
     do { printf("   %-50s ", testName); } while (0)
 #define TPASS() \
-    do { printf("\x1b[32mPASS\x1b[0m\n"); } while (0)
+    do { printf("PASS\n"); } while (0)
 #define TPASSINFO(fmt, ...) \
-    do { printf("\x1b[32mPASS\x1b[0m\n      " fmt "\n", ##__VA_ARGS__); } while (0)
+    do { printf("PASS\n      " fmt "\n", ##__VA_ARGS__); } while (0)
 #define TFAIL() \
     do { \
-      printf("\x1b[31mFAIL\x1b[0m\n"); \
+      printf("FAIL\n"); \
+      [self unloadExtension]; \
       exit(1); \
     } while (0)
 #define TFAILINFO(fmt, ...) \
     do { \
-      printf("\x1b[31mFAIL\x1b[0m\n   -> " fmt "\n\nTest failed.\n\n", ##__VA_ARGS__); \
+      printf("FAIL\n   -> " fmt "\n\nTest failed.\n\n", ##__VA_ARGS__); \
+      [self unloadExtension]; \
       exit(1); \
     } while (0)
 
 @interface SantaKernelTests : NSObject
 @property io_connect_t connection;
-@property int timesSeenLs;
-@property int timesSeenCat;
-@property int timesSeenCp;
 
-@property int testExeIteration;
-@property int timesSeenTestExeIteration;
+// A block that tests can set to handle specific files/binaries.
+// The block should return an action to respond to the kernel with.
+// If no block is specified or no action is returned, the exec will be allowed.
+@property(atomic, copy) santa_action_t (^handlerBlock)(santa_message_t msg);
+
+- (void)unloadDaemon;
+- (void)unloadExtension;
+- (void)loadExtension;
 - (void)runTests;
 @end
 
@@ -94,21 +103,22 @@
 
 /// Call in-kernel function: |kSantaUserClientAllowBinary| or |kSantaUserClientDenyBinary|
 /// passing the |vnodeID|.
-- (void)postToKernelAction:(santa_action_t)action forVnodeID:(uint64_t)vnodeid {
-  if (action == ACTION_RESPOND_ALLOW) {
-    IOConnectCallScalarMethod(self.connection, kSantaUserClientAllowBinary, &vnodeid, 1, 0, 0);
-  } else if (action == ACTION_RESPOND_DENY) {
-    IOConnectCallScalarMethod(self.connection, kSantaUserClientDenyBinary, &vnodeid, 1, 0, 0);
+- (void)postToKernelAction:(santa_action_t)action forVnodeID:(santa_vnode_id_t)vnodeid {
+  if (action == ACTION_RESPOND_DENY) {
+    IOConnectCallStructMethod(self.connection, kSantaUserClientDenyBinary,
+                              &vnodeid, sizeof(vnodeid), 0, 0);
   } else if (action == ACTION_RESPOND_ACK) {
-    IOConnectCallScalarMethod(self.connection, kSantaUserClientAcknowledgeBinary,
-                              &vnodeid, 1, 0, 0);
+    IOConnectCallStructMethod(self.connection, kSantaUserClientAcknowledgeBinary,
+                              &vnodeid, sizeof(vnodeid), 0, 0);
+  } else {
+    IOConnectCallStructMethod(self.connection, kSantaUserClientAllowBinary,
+                              &vnodeid, sizeof(vnodeid), 0, 0);
   }
 }
 
 /// Call in-kernel function: |kSantaUserClientClearCache|
 - (void)flushCache {
-  uint64_t nonRootOnly = 0;
-  IOConnectCallScalarMethod(self.connection, kSantaUserClientClearCache, &nonRootOnly, 1, 0, 0);
+  IOConnectCallScalarMethod(self.connection, kSantaUserClientClearCache, 0, 0, 0, 0);
 }
 
 #pragma mark - Connection Tests
@@ -177,23 +187,16 @@
 /// From then on, monitors the IODataQueue and responds for files specifically used in other tests.
 /// For everything else, allows execution normally to avoid deadlocking the system.
 - (void)beginListening {
-  kern_return_t kr;
-  santa_message_t vdata;
-  UInt32 dataSize;
-  IODataQueueMemory *queueMemory;
-  mach_port_t receivePort;
-
-  mach_vm_address_t address = 0;
-  mach_vm_size_t size = 0;
-
   TSTART("Allocates a notification port");
+  mach_port_t receivePort;
   if (!(receivePort = IODataQueueAllocateNotificationPort())) {
     TFAIL();
   }
   TPASS();
 
   TSTART("Registers the notification port");
-  kr = IOConnectSetNotificationPort(self.connection, QUEUETYPE_DECISION, receivePort, 0);
+  kern_return_t kr = IOConnectSetNotificationPort(
+      self.connection, QUEUETYPE_DECISION, receivePort, 0);
   if (kr != kIOReturnSuccess) {
     mach_port_destroy(mach_task_self(), receivePort);
     TFAILINFO("KR: %d", kr);
@@ -202,6 +205,8 @@
   TPASS();
 
   TSTART("Maps shared memory");
+  mach_vm_address_t address = 0;
+  mach_vm_size_t size = 0;
   kr = IOConnectMapMemory(self.connection, QUEUETYPE_DECISION, mach_task_self(),
                           &address, &size, kIOMapAnywhere);
   if (kr != kIOReturnSuccess) {
@@ -210,85 +215,26 @@
   }
   TPASS();
 
-  // Fetch the SHA-256 of /bin/ed, as we'll be using that for the cache invalidation test.
-  NSString *edSHA = [self sha256ForPath:@"/bin/ed"];
-
-  // Create the RE used for matching testexe's
-  NSString *cwd = [[NSFileManager defaultManager] currentDirectoryPath];
-  NSString *pattern = [cwd stringByAppendingPathComponent:@"testexe\\.(\\d+)"];
-  NSRegularExpression *re = [NSRegularExpression regularExpressionWithPattern:pattern
-                                                                      options:0
-                                                                        error:NULL];
-
   /// Begin listening for events
-  queueMemory = (IODataQueueMemory *)address;
+  IODataQueueMemory *queueMemory = (IODataQueueMemory *)address;
   do {
     while (IODataQueueDataAvailable(queueMemory)) {
-      dataSize = sizeof(vdata);
+      santa_message_t vdata;
+      UInt32 dataSize = sizeof(vdata);
       kr = IODataQueueDequeue(queueMemory, &vdata, &dataSize);
-      if (kr == kIOReturnSuccess) {
-        if (vdata.action != ACTION_REQUEST_BINARY) continue;
-
-        if ([[self sha256ForPath:@(vdata.path)] isEqual:edSHA]) {
-          [self postToKernelAction:ACTION_RESPOND_DENY forVnodeID:vdata.vnode_id];
-        } else if (strncmp("/bin/mv", vdata.path, strlen("/bin/mv")) == 0) {
-          [self postToKernelAction:ACTION_RESPOND_DENY forVnodeID:vdata.vnode_id];
-        } else if (strncmp("/bin/ls", vdata.path, strlen("/bin/ls")) == 0) {
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-          self.timesSeenLs++;
-        } else if (strncmp("/bin/cp", vdata.path, strlen("/bin/cp")) == 0) {
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-          self.timesSeenCp++;
-        } else if (strncmp("/bin/cat", vdata.path, strlen("/bin/cat")) == 0) {
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-          self.timesSeenCat++;
-        } else if (strncmp("/usr/bin/cal", vdata.path, strlen("/usr/bin/cal")) == 0) {
-          static int count = 0;
-          if (count++) TFAILINFO("Large binary should not re-request");
-          [self postToKernelAction:ACTION_RESPOND_ACK forVnodeID:vdata.vnode_id];
-          for (int i = 0; i < 15; ++i) {
-            printf("\033[s");  // save cursor position
-            printf("%i/15", i);
-            sleep(1);
-            printf("\033[u");  // restore cursor position
-          }
-          printf("\033[K\033[u");  // clear line, restore cursor position
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-        } else if (strncmp("/bin/ln", vdata.path, strlen("/bin/ln")) == 0) {
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-
-          TSTART("Sends valid pid/ppid");
-          if (vdata.pid < 1 || vdata.ppid < 1) {
-            TFAIL();
-          }
-          TPASSINFO("Received pid, ppid: %d, %d", vdata.pid, vdata.ppid);
-        } else {
-          NSString *path = @(vdata.path);
-
-          // If current executable is one of our test exe's from handlesLotsOfBinaries,
-          // check that the number has increased.
-          NSArray *matches = [re matchesInString:path
-                                         options:0
-                                           range:NSMakeRange(0, path.length)];
-          if (matches.count == 1 && [matches[0] numberOfRanges] == 2) {
-            NSUInteger count = [[path substringWithRange:[matches[0] rangeAtIndex:1]] intValue];
-            if (count <= self.testExeIteration && count > 0) {
-              self.timesSeenTestExeIteration++;
-              if (self.timesSeenTestExeIteration > 2) {
-                TFAILINFO("Saw same binary several times");
-              }
-            } else {
-              self.timesSeenTestExeIteration = 0;
-              self.testExeIteration = (int)count;
-            }
-          }
-
-          // Allow everything not related to our testing.
-          [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:vdata.vnode_id];
-        }
-      } else {
+      if (kr != kIOReturnSuccess) {
         TFAILINFO("Error receiving data: %d", kr);
+        continue;
       }
+      if (vdata.action != ACTION_REQUEST_BINARY) continue;
+
+      santa_action_t action = ACTION_RESPOND_ALLOW;
+
+      @synchronized(self) {
+        if (self.handlerBlock) action = self.handlerBlock(vdata);
+      }
+
+      [self postToKernelAction:action forVnodeID:vdata.vnode_id];
     }
   } while (IODataQueueWaitForAvailableData(queueMemory, receivePort) == kIOReturnSuccess);
 
@@ -298,12 +244,15 @@
 
 #pragma mark - Functional Tests
 
-/// Tests that blocking works correctly
 - (void)receiveAndBlockTests {
   TSTART("Blocks denied binaries");
 
-  NSTask *ed = [self taskWithPath:@"/bin/ed"];
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    if (strncmp("/bin/ed", msg.path, 7) == 0) return ACTION_RESPOND_DENY;
+    return ACTION_RESPOND_ALLOW;
+  };
 
+  NSTask *ed = [self taskWithPath:@"/bin/ed"];
   @try {
     [ed launch];
     [ed waitUntilExit];
@@ -314,17 +263,20 @@
   }
 }
 
-/// Tests that an allowed binary is cached
 - (void)receiveAndCacheTests {
   TSTART("Permits & caches allowed binaries");
 
-  self.timesSeenLs = 0;
+  __block int timesSeenLs = 0;
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    if (strncmp("/bin/ls", msg.path, 7) == 0) ++timesSeenLs;
+    return ACTION_RESPOND_ALLOW;
+  };
 
   NSTask *ls = [self taskWithPath:@"/bin/ls"];
   [ls launch];
   [ls waitUntilExit];
 
-  if (self.timesSeenLs != 1) {
+  if (timesSeenLs != 1) {
     TFAILINFO("Didn't record first run of ls");
   }
 
@@ -332,25 +284,37 @@
   [ls launch];
   [ls waitUntilExit];
 
-  if (self.timesSeenLs > 1) {
+  if (timesSeenLs > 1) {
     TFAILINFO("Received request for ls a second time");
   }
 
   TPASS();
 }
 
-/// Tests that a write to a cached vnode will invalidate the cached response for that file
 - (void)invalidatesCacheTests {
   TSTART("Invalidates cache for manually closed FDs");
 
-  // Copy the ls binary to a new file
   NSFileManager *fm = [NSFileManager defaultManager];
-  if (![fm copyItemAtPath:@"/bin/pwd" toPath:@"invalidacachetest_tmp" error:nil]) {
+  NSString *target =
+      [[fm currentDirectoryPath] stringByAppendingPathComponent:@"invalidatecachetest"];
+  NSString *edSHA = [self sha256ForPath:@"/bin/ed"];
+
+  __weak __typeof(self) weakSelf = self;
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    __strong __typeof(weakSelf) self = weakSelf;
+    if ([[self sha256ForPath:@(msg.path)] isEqual:edSHA]) {
+      return ACTION_RESPOND_DENY;
+    }
+    return ACTION_RESPOND_ALLOW;
+  };
+
+  // Copy the pwd binary to a new file
+  if (![fm copyItemAtPath:@"/bin/pwd" toPath:target error:nil]) {
     TFAILINFO("Failed to create temp file");
   }
 
   // Launch the new file to put it in the cache
-  NSTask *pwd = [self taskWithPath:@"invalidacachetest_tmp"];
+  NSTask *pwd = [self taskWithPath:target];
   [pwd launch];
   [pwd waitUntilExit];
 
@@ -362,7 +326,7 @@
   // Now replace the contents of the test file (which is cached) with the contents of /bin/ed,
   // which is 'blacklisted' by SHA-256 during the tests.
   FILE *infile = fopen("/bin/ed", "r");
-  FILE *outfile = fopen("invalidacachetest_tmp", "w");
+  FILE *outfile = fopen(target.UTF8String, "w");
   int ch;
   while ((ch = fgetc(infile)) != EOF) {
     fputc(ch, outfile);
@@ -370,13 +334,13 @@
   fclose(infile);
 
   // Now try running the temp file again. If it succeeds, the test failed.
-  NSTask *ed = [self taskWithPath:@"invalidacachetest_tmp"];
+  NSTask *ed = [self taskWithPath:target];
 
   @try {
     [ed launch];
     [ed waitUntilExit];
     TFAILINFO("Launched after write while file open");
-    [fm removeItemAtPath:@"invalidacachetest_tmp" error:nil];
+    [fm removeItemAtPath:target error:nil];
   } @catch (NSException *exception) {
     // This is a pass, but we have more to do.
   }
@@ -385,7 +349,7 @@
   fclose(outfile);
 
   // And try running the temp file again. If it succeeds, the test failed.
-  ed = [self taskWithPath:@"invalidacachetest_tmp"];
+  ed = [self taskWithPath:target];
 
   @try {
     [ed launch];
@@ -394,14 +358,25 @@
   } @catch (NSException *exception) {
     TPASS();
   } @finally {
-    [fm removeItemAtPath:@"invalidacachetest_tmp" error:nil];
+    [fm removeItemAtPath:target error:nil];
   }
 }
 
 - (void)invalidatesCacheAutoCloseTest {
-  TSTART("Invalidates cache for auto-closed FDs");
+  TSTART("Invalidates cache for auto closed FDs");
 
-  // Check invalidations when kernel auto-closes descriptors
+  NSString *edSHA = [self sha256ForPath:@"/bin/ed"];
+
+  __weak __typeof(self) weakSelf = self;
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    __strong __typeof(weakSelf) self = weakSelf;
+    if ([[self sha256ForPath:@(msg.path)] isEqual:edSHA]) {
+      return ACTION_RESPOND_DENY;
+    }
+    return ACTION_RESPOND_ALLOW;
+  };
+
+  // Create temporary file
   NSFileManager *fm = [NSFileManager defaultManager];
   if (![fm copyItemAtPath:@"/bin/pwd" toPath:@"invalidacachetest_tmp" error:nil]) {
     TFAILINFO("Failed to create temp file");
@@ -411,13 +386,11 @@
   NSTask *pwd = [self taskWithPath:@"invalidacachetest_tmp"];
   [pwd launch];
   [pwd waitUntilExit];
-
-  // Exit if this fails with a useful message.
   if ([pwd terminationStatus] != 0) {
     TFAILINFO("Second launch of test binary failed");
   }
 
-  // Replace file contents
+  // Replace file contents using dd, which doesn't close FDs
   NSDictionary *attrs = [fm attributesOfItemAtPath:@"/bin/ed" error:NULL];
   NSTask *dd = [self taskWithPath:@"/bin/dd"];
   dd.arguments = @[ @"if=/bin/ed",
@@ -441,17 +414,20 @@
   }
 }
 
-/// Tests the clear cache function works correctly
 - (void)clearCacheTests {
   TSTART("Can clear cache");
 
-  self.timesSeenCat = 0;
+  __block int timesSeenCat = 0;
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    if (strncmp("/bin/cat", msg.path, 8) == 0) ++timesSeenCat;
+    return ACTION_RESPOND_ALLOW;
+  };
 
   NSTask *cat = [self taskWithPath:@"/bin/cat"];
   [cat launch];
   [cat waitUntilExit];
 
-  if (self.timesSeenCat != 1) {
+  if (timesSeenCat != 1) {
     TFAILINFO("Didn't record first run of cat");
   }
 
@@ -461,23 +437,27 @@
   [cat launch];
   [cat waitUntilExit];
 
-  if (self.timesSeenCat != 2) {
+  if (timesSeenCat != 2) {
     TFAIL();
   }
 
   TPASS();
 }
 
-/// Tests that the kernel still denies blocked binaries even if launched while traced
 - (void)blocksDeniedTracedBinaries {
   TSTART("Denies blocked processes running while traced");
 
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    if (strncmp("/bin/mv", msg.path, 7) == 0) return ACTION_RESPOND_DENY;
+    return ACTION_RESPOND_ALLOW;
+  };
+
   pid_t pid = fork();
   if (pid < 0) {
     TFAILINFO("Failed to fork");
   } else if (pid > 0) {
     int status;
-    waitpid(pid, &status, 0);
+    while (waitpid(pid, &status, 0) != pid); // handle EINTR
     if (WIFEXITED(status) && WEXITSTATUS(status) == EPERM) {
       TPASS();
     } else if (WIFSTOPPED(status)) {
@@ -494,44 +474,12 @@
   }
 }
 
-/// Tests that the kernel can handle _lots_ of executions.
-- (void)handlesLotsOfBinaries {
-  TSTART("Handles lots of binaries");
-
-  const int LIMIT = 12000;
-
-  for (int i = 0; i < LIMIT; ++i) {
-    printf("\033[s");  // save cursor position
-
-    printf("%d/%i", i + 1, LIMIT);
-
-    NSString *fname = [@"testexe" stringByAppendingFormat:@".%i", i];
-    [[NSFileManager defaultManager] copyItemAtPath:@"/bin/hostname" toPath:fname error:NULL];
-
-    @try {
-      NSTask *testexec = [self taskWithPath:fname];
-      [testexec launch];
-      [testexec waitUntilExit];
-    } @catch (NSException *e) {
-      TFAILINFO("Failed to launch");
-    }
-
-    unlink([fname UTF8String]);
-    printf("\033[u");  // restore cursor position
-  }
-  printf("\033[K\033[u");  // clear line, restore cursor position
-
-  TPASS();
-}
-
 - (void)testCachePerformance {
-  TSTART("Test cache performance");
+  TSTART("Test cache performance...");
 
   // Execute echo 100 times, saving the time taken for each run
   std::vector<double> times;
   for (int i = 0; i < 100; ++i) {
-    printf("\033[s");  // save cursor position
-    printf("%d/%d", i + 1, 100);
     NSTask *t = [[NSTask alloc] init];
     t.launchPath = @"/bin/echo";
     t.standardOutput = [NSPipe pipe];
@@ -542,11 +490,8 @@
     auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
 
     if (i > 5) times.push_back(duration);
-    printf("\033[u");  // restore cursor position
   }
 
-  printf("\033[K\033[u");  // clear line, restore cursor position
-
   // Sort and remove first 10 and last 10 entries.
   std::sort(times.begin(), times.end());
   times.erase(times.begin(), times.begin()+10);
@@ -570,12 +515,31 @@
 }
 
 - (void)testLargeBinary {
-  TSTART("Handles large binary");
+  TSTART("Handles large binary...");
+
+  __block int calCount = 0;
+  __weak __typeof(self) weakSelf = self;
+  self.handlerBlock = ^santa_action_t(santa_message_t msg) {
+    __strong __typeof(weakSelf) self = weakSelf;
+    if (strncmp("/usr/bin/cal", msg.path, 12) == 0) {
+      if (calCount++) TFAILINFO("Large binary should not re-request");
+      dispatch_after(dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC),
+                     dispatch_get_global_queue(0, 0), ^{
+        [self postToKernelAction:ACTION_RESPOND_ALLOW forVnodeID:msg.vnode_id];
+      });
+      return ACTION_RESPOND_ACK;
+    }
+    return ACTION_RESPOND_ALLOW;
+  };
 
   @try {
     NSTask *testexec = [self taskWithPath:@"/usr/bin/cal"];
     [testexec launch];
-    [testexec waitUntilExit];
+    int sleepCount = 0;
+    while ([testexec isRunning]) {
+      sleep(1);
+      if (++sleepCount > 5) TFAILINFO("Took longer than expected to start/stop");
+    }
   } @catch (NSException *e) {
     TFAILINFO("Failed to launch");
   }
@@ -585,8 +549,57 @@
 
 #pragma mark - Main
 
+- (void)unloadDaemon {
+  NSTask *t = [[NSTask alloc] init];
+  t.launchPath = @"/bin/launchctl";
+  t.arguments = @[ @"remove", @"com.google.santad" ];
+  t.standardOutput = t.standardError = [NSPipe pipe];
+  [t launch];
+  [t waitUntilExit];
+}
+
+- (void)unloadExtension {
+  // Don't check the status of this, the kext may not be loaded..
+  OSStatus ret = KextManagerUnloadKextWithIdentifier(CFSTR("com.google.santa-driver"));
+  if (ret != kOSReturnSuccess && ret != kOSKextReturnNotFound) {
+    NSLog(@"Failed to unload extension: 0x%X", ret);
+  }
+}
+
+- (void)loadExtension {
+  TSTART("Loads extension");
+
+  NSError *error;
+  NSFileManager *fm = [NSFileManager defaultManager];
+
+  NSString *src = [[fm currentDirectoryPath] stringByAppendingPathComponent:@"santa-driver.kext"];
+  NSString *dest = [NSTemporaryDirectory() stringByAppendingPathComponent:@"santa-driver.kext"];
+  [fm removeItemAtPath:dest error:NULL]; // ensure dest is free
+  if (![fm copyItemAtPath:src toPath:dest error:&error] || error) {
+    TFAILINFO("Failed to copy kext: %s", error.description.UTF8String);
+  }
+
+  NSDictionary *attrs = @{
+      NSFileOwnerAccountName : @"root",
+      NSFileGroupOwnerAccountName : @"wheel",
+      NSFilePosixPermissions : @0755
+  };
+
+  [fm setAttributes:attrs ofItemAtPath:dest error:NULL];
+  for (NSString *path in [fm enumeratorAtPath:dest]) {
+    [fm setAttributes:attrs ofItemAtPath:[dest stringByAppendingPathComponent:path] error:NULL];
+  }
+
+  NSURL *destURL = [NSURL fileURLWithPath:dest];
+  OSStatus ret = KextManagerLoadKextWithURL((__bridge CFURLRef)destURL, NULL);
+  if (ret != kOSReturnSuccess) {
+    TFAILINFO("Failed to load kext: 0x%X", ret);
+  }
+  usleep(50000);
+  TPASS();
+}
+
 - (void)runTests {
-  printf("\nSanta Kernel Tests\n==================\n");
   printf("-> Connection tests:\n");
 
   // Test that connection can be established
@@ -598,19 +611,18 @@
 
   // Wait for driver to finish getting ready
   sleep(1);
-  printf("\n-> Functional tests:\033[m\n");
 
+  printf("\n-> Functional tests:\n");
   [self receiveAndBlockTests];
   [self receiveAndCacheTests];
   [self invalidatesCacheTests];
   [self invalidatesCacheAutoCloseTest];
   [self clearCacheTests];
   [self blocksDeniedTracedBinaries];
-
-  printf("\n-> Performance tests:\033[m\n");
-  [self testCachePerformance];
   [self testLargeBinary];
-  [self handlesLotsOfBinaries];
+
+  printf("\n-> Performance tests:\n");
+  [self testCachePerformance];
 
   printf("\nAll tests passed.\n\n");
 }
@@ -626,8 +638,18 @@ int main(int argc, const char *argv[]) {
       exit(1);
     }
 
+    chdir([[[NSBundle mainBundle] bundlePath] UTF8String]);
+
     SantaKernelTests *skt = [[SantaKernelTests alloc] init];
+    printf("\nSanta Kernel Tests\n==================\n\n");
+    printf("-> Loading tests:\n");
+    [skt unloadDaemon];
+    [skt unloadExtension];
+    [skt loadExtension];
+    printf("\n");
+
     [skt runTests];
+    [skt unloadExtension];
   }
   return 0;
 }

Tests/LogicTests/SNTExecutionControllerTest.m

@@ -80,11 +80,15 @@
   santa_message_t message = {0};
   message.pid = 12;
   message.ppid = 1;
-  message.vnode_id = 1234;
+  message.vnode_id = [self getVnodeId];
   strncpy(message.path, "/a/file", 7);
   return message;
 }
 
+- (santa_vnode_id_t)getVnodeId {
+  return (santa_vnode_id_t){.fsid = 1234, .fileid = 5678};
+}
+
 - (void)tearDown {
   [self.mockFileInfo stopMocking];
   [self.mockCodesignChecker stopMocking];
@@ -107,7 +111,7 @@
   [self.sut validateBinaryWithMessage:[self getMessage]];
 
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_ALLOW
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testBinaryBlacklistRule {
@@ -122,7 +126,7 @@
   [self.sut validateBinaryWithMessage:[self getMessage]];
 
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_DENY
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testCertificateWhitelistRule {
@@ -140,7 +144,7 @@
   [self.sut validateBinaryWithMessage:[self getMessage]];
 
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_ALLOW
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testCertificateBlacklistRule {
@@ -158,7 +162,7 @@
   [self.sut validateBinaryWithMessage:[self getMessage]];
 
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_DENY
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testDefaultDecision {
@@ -168,12 +172,12 @@
   OCMExpect([self.mockConfigurator clientMode]).andReturn(SNTClientModeMonitor);
   [self.sut validateBinaryWithMessage:[self getMessage]];
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_ALLOW
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 
   OCMExpect([self.mockConfigurator clientMode]).andReturn(SNTClientModeLockdown);
   [self.sut validateBinaryWithMessage:[self getMessage]];
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_DENY
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testOutOfScope {
@@ -182,13 +186,13 @@
   OCMStub([self.mockConfigurator clientMode]).andReturn(SNTClientModeLockdown);
   [self.sut validateBinaryWithMessage:[self getMessage]];
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_ALLOW
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testMissingShasum {
   [self.sut validateBinaryWithMessage:[self getMessage]];
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_ALLOW
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 - (void)testPageZero {
@@ -197,7 +201,7 @@
 
   [self.sut validateBinaryWithMessage:[self getMessage]];
   OCMVerify([self.mockDriverManager postToKernelAction:ACTION_RESPOND_DENY
-                                            forVnodeID:1234]);
+                                            forVnodeID:[self getVnodeId]]);
 }
 
 @end

Tests/LogicTests/SantaCacheTest.mm

@@ -14,7 +14,9 @@
 
 #import <XCTest/XCTest.h>
 
+#include <numeric>
 #include <string>
+#include <vector>
 
 #include "SantaCache.h"
 
@@ -28,71 +30,113 @@
 }
 
 - (void)testSetAndGet {
-  auto sut = new SantaCache<uint64_t>();
+  auto sut = SantaCache<uint64_t, uint64_t>();
 
-  sut->set(72057611258548992llu, 10000192);
-  XCTAssertEqual(sut->get(72057611258548992llu), 10000192);
-
-  delete sut;
+  sut.set(72057611258548992llu, 10000192);
+  XCTAssertEqual(sut.get(72057611258548992llu), 10000192);
 }
 
 - (void)testCacheRemove {
-  auto sut = new SantaCache<uint64_t>();
+  auto sut = SantaCache<uint64_t, uint64_t>();
 
-  sut->set(0xDEADBEEF, 42);
-  sut->remove(0xDEADBEEF);
+  sut.set(0xDEADBEEF, 42);
+  sut.remove(0xDEADBEEF);
 
-  XCTAssertEqual(sut->get(0xDEADBEEF), 0);
-
-  delete sut;
+  XCTAssertEqual(sut.get(0xDEADBEEF), 0);
 }
 
 - (void)testBucketGrowCopy {
-  auto sut = new SantaCache<uint64_t>();
+  auto sut = SantaCache<uint64_t, uint64_t>();
 
-  sut->set(386, 42);
-  sut->set(2434, 42);
+  sut.set(386, 42);
+  sut.set(2434, 42);
 
-  XCTAssertEqual(sut->get(386), 42);
-  XCTAssertEqual(sut->get(2434), 42);
-
-  delete sut;
+  XCTAssertEqual(sut.get(386), 42);
+  XCTAssertEqual(sut.get(2434), 42);
 }
 
 - (void)testBucketShrinkCopy {
-  auto sut = new SantaCache<uint64_t>(100, 1);
+  auto sut = SantaCache<uint64_t, uint64_t>(100, 1);
 
-  sut->set(386, 42);
-  sut->set(2434, 42);
-  sut->set(4482, 42);
+  sut.set(386, 42);
+  sut.set(2434, 42);
+  sut.set(4482, 42);
 
-  sut->remove(2434);
+  sut.remove(2434);
 
-  XCTAssertEqual(sut->get(386), 42);
-  XCTAssertEqual(sut->get(2434), 0);
-  XCTAssertEqual(sut->get(4482), 42);
-
-  delete sut;
+  XCTAssertEqual(sut.get(386), 42);
+  XCTAssertEqual(sut.get(2434), 0);
+  XCTAssertEqual(sut.get(4482), 42);
 }
 
 - (void)testCacheResetAtLimit {
-  auto sut = new SantaCache<uint64_t>(5);
+  auto sut = SantaCache<uint64_t, uint64_t>(5);
   
-  sut->set(1, 42);
-  sut->set(2, 42);
-  sut->set(3, 42);
-  sut->set(4, 42);
-  sut->set(5, 42);
-  XCTAssertEqual(sut->get(3), 42);
-  sut->set(6, 42);
-  XCTAssertEqual(sut->get(3), 0);
-  XCTAssertEqual(sut->get(6), 42);
+  sut.set(1, 42);
+  sut.set(2, 42);
+  sut.set(3, 42);
+  sut.set(4, 42);
+  sut.set(5, 42);
+  XCTAssertEqual(sut.get(3), 42);
+  sut.set(6, 42);
+  XCTAssertEqual(sut.get(3), 0);
+  XCTAssertEqual(sut.get(6), 42);
+}
 
-  delete sut;
+// Helper to test bucket distributions for uint64_t/uint64_t combinations.
+- (void)distributionTestHelper:(SantaCache<uint64_t, uint64_t> *)sut bucketRatio:(int)br {
+  uint16_t count[512];
+  uint16_t array_size = 512;
+  uint64_t start_bucket = 0;
+  std::vector<uint16_t> per_bucket;
+  do {
+    sut->bucket_counts(count, &array_size, &start_bucket);
+    for (int i = 0; i < array_size; ++i) {
+      per_bucket.push_back(count[i]);
+    }
+  } while (start_bucket > 0);
+
+  // Calculate mean
+  double mean = std::accumulate(per_bucket.begin(), per_bucket.end(), 0.0) / per_bucket.size();
+  XCTAssertLessThanOrEqual(mean, br, @"Mean per-bucket count is greater than %d", br);
+
+  // Calculate stdev
+  double accum = 0.0;
+  std::for_each(per_bucket.begin(), per_bucket.end(), [&](const double d) {
+    accum += (d - mean) * (d - mean);
+  });
+  double stddev = sqrt(accum / (per_bucket.size() - 1));
+  double maxStdDev = (double)br / 2;
+  XCTAssertLessThanOrEqual(stddev, maxStdDev,
+                           @"Standard deviation between buckets is greater than %f", maxStdDev);
+}
+
+- (void)testDistributionRandomKeys {
+  const int bucket_ratio = 5;
+  auto sut = new SantaCache<uint64_t, uint64_t>(5000, bucket_ratio);
+
+  // Fill the cache with random keys, all set to 1.
+  for (int i = 0; i < 4000; ++i) {
+    sut->set((uint64_t)arc4random() << 32 | arc4random(), 1);
+  }
+
+  [self distributionTestHelper:sut bucketRatio:bucket_ratio];
+}
+
+- (void)testDistributionMontonicKeys {
+  const int bucket_ratio = 5;
+  auto sut = new SantaCache<uint64_t, uint64_t>(5000, bucket_ratio);
+
+  // Fill the cache with monotonic keys, all set to 1.
+  for (int i = 0; i < 4000; ++i) {
+    sut->set(i, 1);
+  }
+
+  [self distributionTestHelper:sut bucketRatio:bucket_ratio];
 }
 
 - (void)testThreading {
-  auto sut = new SantaCache<uint64_t>();
+  auto sut = new SantaCache<uint64_t, uint64_t>();
 
   for (int x = 0; x < 200; ++x) {
     dispatch_group_t group = dispatch_group_create();
@@ -120,57 +164,105 @@
 }
 
 - (void)testCount {
-  auto sut = new SantaCache<uint64_t>();
+  auto sut = SantaCache<uint64_t, int>();
 
-  XCTAssertEqual(sut->count(), 0);
+  XCTAssertEqual(sut.count(), 0);
 
-  sut->set(4012, 42);
-  sut->set(42, 0);  
-  sut->set(0x8BADF00D, 40010);
+  sut.set(4012, 42);
+  sut.set(42, 0);
+  sut.set(0x8BADF00D, 40010);
 
-  XCTAssertEqual(sut->count(), 2);
+  XCTAssertEqual(sut.count(), 2);
+}
 
-  delete sut;
+- (void)testDoubles {
+  auto sut = SantaCache<double, double>();
+
+  sut.set(3.14, 2.718281);
+  sut.set(1.41429, 2.5029);
+  sut.set(4.6692, 1.2020569);
+  sut.set(1.61803398, 0.57721);
+
+  XCTAssertEqual(sut.count(), 4);
+  XCTAssertEqual(sut.get(3.14), 2.718281);
+  XCTAssertEqual(sut.get(1.41429), 2.5029);
+  XCTAssertEqual(sut.get(4.6692), 1.2020569);
+  XCTAssertEqual(sut.get(1.61803398), 0.57721);
+
+  XCTAssertEqual(sut.get(5.5555), 0);
+  XCTAssertEqual(sut.get(3.1459124), 0);
+}
+
+template<> uint64_t SantaCacheHasher<std::string>(std::string const& s) {
+  return std::hash<std::string>{}(s);
 }
 
 - (void)testStrings {
-  auto sut = new SantaCache<std::string>();
+  auto sut = SantaCache<std::string, std::string>();
 
-  sut->set(1, "deadbeef");
-  sut->set(2, "feedface");
+  std::string s1 = "foo";
+  std::string s2 = "bar";
 
-  XCTAssertEqual(sut->count(), 2);
-  XCTAssertEqual(sut->get(1), "deadbeef");
-  XCTAssertEqual(sut->get(2), "feedface");
+  sut.set(s1, "deadbeef");
+  sut.set(s2, "feedface");
 
-  sut->remove(2);
+  XCTAssertEqual(sut.count(), 2);
+  XCTAssertEqual(sut.get(s1), "deadbeef");
+  XCTAssertEqual(sut.get(s2), "feedface");
 
-  XCTAssertTrue(sut->get(2).empty());
+  sut.remove(s2);
 
-  delete sut;
+  XCTAssertTrue(sut.get(s2).empty());
 }
 
 - (void)testCompareAndSwap {
-  auto sut = new SantaCache<uint64_t>(100, 2);
+  auto sut = SantaCache<uint64_t, 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, 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, 0);
+  XCTAssertEqual(sut.get(1), 0);
 
-  sut->set(1, 42, 1);
-  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, 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);
+  sut.set(1, 0, 666);
+  XCTAssertEqual(sut.get(1), 42);
+  sut.set(1, 0, 42);
+  XCTAssertEqual(sut.get(1), 0);
+}
+
+struct S {
+  uint64_t first_val;
+  uint64_t second_val;
+
+  bool operator==(const S& rhs) {
+    return first_val == rhs.first_val && second_val == rhs.second_val;
+  }
+};
+template<> uint64_t SantaCacheHasher<S>(S const& s) {
+  return SantaCacheHasher<uint64_t>(s.first_val) ^ (SantaCacheHasher<uint64_t>(s.second_val) << 1);
+}
+
+- (void)testStructKeys {
+  auto sut =  SantaCache<S, uint64_t>(10, 2);
+
+  S s1 = {1024, 2048};
+  S s2 = {4096, 8192};
+  S s3 = {16384, 32768};
+  sut.set(s1, 10);
+  sut.set(s2, 20);
+  sut.set(s3, 30);
+
+  XCTAssertEqual(sut.get(s1), 10);
+  XCTAssertEqual(sut.get(s2), 20);
+  XCTAssertEqual(sut.get(s3), 30);
 }
 
 @end