Update SNTPolicyProcessor to use a map (#1304)

* Update SNTPolicyProcessor to use a map instead of a giant switch statement

Update SNTPolicyProcessor to use a map instead of a giant switch statement.

Add unit tests for the method that sets SNTCachedDecision values.

* Remove unneccessary OCMock dep in BUILD file.

* Fix typo in method signature.

* Incorporate review feedback.

* Upper case UpdateCachedDecisionSigningInfo

* Update SNTPolicyProcessor.h

Co-authored-by: Russell Hancox <russellhancox@users.noreply.github.com>

* Update SNTPolicyProcessor.mm

Co-authored-by: Russell Hancox <russellhancox@users.noreply.github.com>

* Fix typo

* Fix linter issues.

* Fixed up more linter issues.

---------

Co-authored-by: Russell Hancox <russellhancox@users.noreply.github.com>

Source/santad/BUILD

@@ -193,7 +193,7 @@ objc_library(
 
 objc_library(
     name = "SNTPolicyProcessor",
-    srcs = ["SNTPolicyProcessor.m"],
+    srcs = ["SNTPolicyProcessor.mm"],
     hdrs = ["SNTPolicyProcessor.h"],
     deps = [
         ":SNTRuleTable",
@@ -209,6 +209,18 @@ objc_library(
         "@MOLCertificate",
         "@MOLCodesignChecker",
         "@MOLXPCConnection",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+santa_unit_test(
+    name = "SNTPolicyProcessorTest",
+    srcs = ["SNTPolicyProcessorTest.mm"],
+    deps = [
+        ":SNTPolicyProcessor",
+        "//Source/common:SNTConfigurator",
+        "//Source/common:SNTRule",
+        "//Source/common:TestUtils",
     ],
 )
 
@@ -1417,6 +1429,7 @@ test_suite(
         ":SNTEndpointSecurityTamperResistanceTest",
         ":SNTEventTableTest",
         ":SNTExecutionControllerTest",
+        ":SNTPolicyProcessorTest",
         ":SNTRuleTableTest",
         ":SantadTest",
         ":WatchItemsTest",

Source/santad/SNTPolicyProcessor.h

@@ -17,6 +17,7 @@
 #import <MOLCertificate/MOLCertificate.h>
 
 #import "Source/common/SNTCommonEnums.h"
+#import "Source/common/SNTRule.h"
 #import "Source/common/SNTRuleIdentifiers.h"
 
 @class MOLCodesignChecker;
@@ -60,4 +61,12 @@
 - (nonnull SNTCachedDecision *)decisionForFilePath:(nonnull NSString *)filePath
                                        identifiers:(nonnull SNTRuleIdentifiers *)identifiers;
 
+///
+/// Updates a decision for a given file and agent configuration.
+///
+/// Returns YES if the decision requires no futher processing NO otherwise.
+- (BOOL)decision:(nonnull SNTCachedDecision *)cd
+              forRule:(nonnull SNTRule *)rule
+  withTransitiveRules:(BOOL)transitive;
+
 @end

Source/santad/SNTPolicyProcessor.m → Source/santad/SNTPolicyProcessor.mm

@@ -27,6 +27,7 @@
 #import "Source/common/SNTLogging.h"
 #import "Source/common/SNTRule.h"
 #import "Source/santad/DataLayer/SNTRuleTable.h"
+#include "absl/container/flat_hash_map.h"
 
 @interface SNTPolicyProcessor ()
 @property SNTRuleTable *ruleTable;
@@ -44,6 +45,118 @@ - (instancetype)initWithRuleTable:(SNTRuleTable *)ruleTable {
   return self;
 }
 
+// This method applies the rules to the cached decision object.
+//
+// It returns YES if the decision was made, NO if the decision was not made.
+- (BOOL)decision:(SNTCachedDecision *)cd
+              forRule:(SNTRule *)rule
+  withTransitiveRules:(BOOL)enableTransitiveRules {
+  static const auto decisions =
+    absl::flat_hash_map<std::pair<SNTRuleType, SNTRuleState>, SNTEventState>{
+      {{SNTRuleTypeCDHash, SNTRuleStateAllow}, SNTEventStateAllowCDHash},
+      {{SNTRuleTypeCDHash, SNTRuleStateAllowCompiler}, SNTEventStateAllowCompiler},
+      {{SNTRuleTypeCDHash, SNTRuleStateBlock}, SNTEventStateBlockCDHash},
+      {{SNTRuleTypeCDHash, SNTRuleStateSilentBlock}, SNTEventStateBlockCDHash},
+      {{SNTRuleTypeBinary, SNTRuleStateAllow}, SNTEventStateAllowBinary},
+      {{SNTRuleTypeBinary, SNTRuleStateAllowTransitive}, SNTEventStateAllowTransitive},
+      {{SNTRuleTypeBinary, SNTRuleStateAllowCompiler}, SNTEventStateAllowCompiler},
+      {{SNTRuleTypeBinary, SNTRuleStateSilentBlock}, SNTEventStateBlockBinary},
+      {{SNTRuleTypeBinary, SNTRuleStateBlock}, SNTEventStateBlockBinary},
+      {{SNTRuleTypeSigningID, SNTRuleStateAllow}, SNTEventStateAllowSigningID},
+      {{SNTRuleTypeSigningID, SNTRuleStateAllowCompiler}, SNTEventStateAllowCompiler},
+      {{SNTRuleTypeSigningID, SNTRuleStateSilentBlock}, SNTEventStateBlockSigningID},
+      {{SNTRuleTypeSigningID, SNTRuleStateBlock}, SNTEventStateBlockSigningID},
+      {{SNTRuleTypeCertificate, SNTRuleStateAllow}, SNTEventStateAllowCertificate},
+      {{SNTRuleTypeCertificate, SNTRuleStateSilentBlock}, SNTEventStateBlockCertificate},
+      {{SNTRuleTypeCertificate, SNTRuleStateBlock}, SNTEventStateBlockCertificate},
+      {{SNTRuleTypeTeamID, SNTRuleStateAllow}, SNTEventStateAllowTeamID},
+      {{SNTRuleTypeTeamID, SNTRuleStateSilentBlock}, SNTEventStateBlockTeamID},
+      {{SNTRuleTypeTeamID, SNTRuleStateBlock}, SNTEventStateBlockTeamID},
+    };
+
+  auto iterator = decisions.find(std::pair<SNTRuleType, SNTRuleState>{rule.type, rule.state});
+  if (iterator != decisions.end()) {
+    cd.decision = iterator->second;
+  } else {
+    // If we have an invalid state combination then either we have stale data in
+    // the database or a programming error. We treat this as if the
+    // corresponding rule was not found.
+    LOGE(@"Invalid rule type/state combination %ld/%ld", rule.type, rule.state);
+    return NO;
+  }
+
+  switch (rule.state) {
+    case SNTRuleStateSilentBlock: cd.silentBlock = YES; break;
+    case SNTRuleStateAllowCompiler:
+      if (!enableTransitiveRules) {
+        switch (rule.type) {
+          case SNTRuleTypeCDHash: cd.decision = SNTEventStateAllowCDHash; break;
+          case SNTRuleTypeBinary: cd.decision = SNTEventStateAllowBinary; break;
+          case SNTRuleTypeSigningID: cd.decision = SNTEventStateAllowSigningID; break;
+          default:
+            // Programming error. Something's marked as a compiler that shouldn't
+            // be.
+            LOGE(@"Invalid compiler rule type %ld", rule.type);
+            [NSException
+               raise:@"Invalid compiler rule type"
+              format:@"decision:forRule:withTransitiveRules: Unexpected compiler rule type: %ld",
+                     rule.type];
+            break;
+        }
+      }
+      break;
+    case SNTRuleStateAllowTransitive:
+      // If transitive rules are disabled, then we treat
+      // SNTRuleStateAllowTransitive rules as if a matching rule was not found
+      // and set the state to unknown. Otherwise the decision map will have already set
+      // the EventState to SNTEventStateAllowTransitive.
+      if (!enableTransitiveRules) {
+        cd.decision = SNTEventStateUnknown;
+        return NO;
+      }
+      break;
+    default:
+      // If its not one of the special cases above, we don't need to do anything.
+      break;
+  }
+
+  // We know we have a match so apply the custom messages
+  cd.customMsg = rule.customMsg;
+  cd.customURL = rule.customURL;
+
+  return YES;
+}
+
+static void UpdateCachedDecisionSigningInfo(
+  SNTCachedDecision *cd, MOLCodesignChecker *csInfo,
+  NSDictionary *_Nullable (^entitlementsFilterCallback)(NSDictionary *_Nullable entitlements)) {
+  cd.certSHA256 = csInfo.leafCertificate.SHA256;
+  cd.certCommonName = csInfo.leafCertificate.commonName;
+  cd.certChain = csInfo.certificates;
+  // Check if we need to get teamID from code signing.
+  if (!cd.teamID) {
+    cd.teamID = csInfo.teamID;
+  }
+
+  // Ensure that if no teamID exists that the signing info confirms it is a
+  // platform binary. If not, remove the signingID.
+  if (!cd.teamID && cd.signingID) {
+    if (!csInfo.platformBinary) {
+      cd.signingID = nil;
+    }
+  }
+
+  NSDictionary *entitlements = csInfo.entitlements;
+
+  if (entitlementsFilterCallback) {
+    cd.entitlements = entitlementsFilterCallback(entitlements);
+    cd.entitlementsFiltered = (cd.entitlements.count != entitlements.count);
+  } else {
+    cd.entitlements = [entitlements sntDeepCopy];
+    cd.entitlementsFiltered = NO;
+  }
+}
+
 - (nonnull SNTCachedDecision *)decisionForFileInfo:(nonnull SNTFileInfo *)fileInfo
                                             cdhash:(nullable NSString *)cdhash
                                         fileSHA256:(nullable NSString *)fileSHA256
@@ -54,23 +167,27 @@ - (nonnull SNTCachedDecision *)decisionForFileInfo:(nonnull SNTFileInfo *)fileIn
                         entitlementsFilterCallback:
                           (NSDictionary *_Nullable (^_Nullable)(
                             NSDictionary *_Nullable entitlements))entitlementsFilterCallback {
-  SNTCachedDecision *cd = [[SNTCachedDecision alloc] init];
-  cd.cdhash = cdhash;
-  cd.sha256 = fileSHA256 ?: fileInfo.SHA256;
-  cd.teamID = teamID;
-  cd.signingID = signingID;
-
+  // Check the hash before allocating a SNTCachedDecision.
+  NSString *fileHash = fileSHA256 ?: fileInfo.SHA256;
   SNTClientMode mode = [self.configurator clientMode];
-  cd.decisionClientMode = mode;
 
   // If the binary is a critical system binary, don't check its signature.
   // The binary was validated at startup when the rule table was initialized.
-  SNTCachedDecision *systemCd = self.ruleTable.criticalSystemBinaries[cd.sha256];
+  SNTCachedDecision *systemCd = self.ruleTable.criticalSystemBinaries[fileHash];
   if (systemCd) {
     systemCd.decisionClientMode = mode;
     return systemCd;
   }
 
+  // Allocate a new cached decision for the execution.
+  SNTCachedDecision *cd = [[SNTCachedDecision alloc] init];
+  cd.cdhash = cdhash;
+  cd.sha256 = fileHash;
+  cd.teamID = teamID;
+  cd.signingID = signingID;
+  cd.decisionClientMode = mode;
+  cd.quarantineURL = fileInfo.quarantineDataURL;
+
   NSError *csInfoError;
   if (certificateSHA256.length) {
     cd.certSHA256 = certificateSHA256;
@@ -87,31 +204,9 @@ - (nonnull SNTCachedDecision *)decisionForFileInfo:(nonnull SNTFileInfo *)fileIn
       cd.signingID = nil;
       cd.cdhash = nil;
     } else {
-      cd.certSHA256 = csInfo.leafCertificate.SHA256;
-      cd.certCommonName = csInfo.leafCertificate.commonName;
-      cd.certChain = csInfo.certificates;
-      cd.teamID = teamID ?: csInfo.teamID;
-
-      // Ensure that if no teamID exists that the signing info confirms it is a
-      // platform binary. If not, remove the signingID.
-      if (!cd.teamID && cd.signingID) {
-        if (!csInfo.platformBinary) {
-          cd.signingID = nil;
-        }
-      }
-
-      NSDictionary *entitlements = csInfo.entitlements;
-
-      if (entitlementsFilterCallback) {
-        cd.entitlements = entitlementsFilterCallback(entitlements);
-        cd.entitlementsFiltered = (cd.entitlements.count == entitlements.count);
-      } else {
-        cd.entitlements = [entitlements sntDeepCopy];
-        cd.entitlementsFiltered = NO;
-      }
+      UpdateCachedDecisionSigningInfo(cd, csInfo, entitlementsFilterCallback);
     }
   }
-  cd.quarantineURL = fileInfo.quarantineDataURL;
 
   // Do not evaluate TeamID/SigningID rules for dev-signed code based on the
   // assumption that orgs are generally more relaxed about dev signed cert
@@ -133,116 +228,11 @@ - (nonnull SNTCachedDecision *)decisionForFileInfo:(nonnull SNTFileInfo *)fileIn
                                                                 .certificateSHA256 = cd.certSHA256,
                                                                 .teamID = cd.teamID}];
   if (rule) {
-    switch (rule.type) {
-      case SNTRuleTypeCDHash:
-        switch (rule.state) {
-          case SNTRuleStateAllow: cd.decision = SNTEventStateAllowCDHash; return cd;
-          case SNTRuleStateAllowCompiler:
-            // If transitive rules are enabled, then SNTRuleStateAllowListCompiler rules
-            // become SNTEventStateAllowCompiler decisions.  Otherwise we treat the rule as if
-            // it were SNTRuleStateAllowCDHash.
-            if ([self.configurator enableTransitiveRules]) {
-              cd.decision = SNTEventStateAllowCompiler;
-            } else {
-              cd.decision = SNTEventStateAllowCDHash;
-            }
-            return cd;
-          case SNTRuleStateSilentBlock:
-            cd.silentBlock = YES;
-            // intentional fallthrough
-          case SNTRuleStateBlock:
-            cd.customMsg = rule.customMsg;
-            cd.customURL = rule.customURL;
-            cd.decision = SNTEventStateBlockCDHash;
-            return cd;
-          default: break;
-        }
-      case SNTRuleTypeBinary:
-        switch (rule.state) {
-          case SNTRuleStateAllow: cd.decision = SNTEventStateAllowBinary; return cd;
-          case SNTRuleStateSilentBlock: cd.silentBlock = YES;
-          case SNTRuleStateBlock:
-            cd.customMsg = rule.customMsg;
-            cd.customURL = rule.customURL;
-            cd.decision = SNTEventStateBlockBinary;
-            return cd;
-          case SNTRuleStateAllowCompiler:
-            // If transitive rules are enabled, then SNTRuleStateAllowListCompiler rules
-            // become SNTEventStateAllowCompiler decisions.  Otherwise we treat the rule as if
-            // it were SNTRuleStateAllow.
-            if ([self.configurator enableTransitiveRules]) {
-              cd.decision = SNTEventStateAllowCompiler;
-            } else {
-              cd.decision = SNTEventStateAllowBinary;
-            }
-            return cd;
-          case SNTRuleStateAllowTransitive:
-            // If transitive rules are enabled, then SNTRuleStateAllowTransitive
-            // rules become SNTEventStateAllowTransitive decisions.  Otherwise, we treat the
-            // rule as if it were SNTRuleStateUnknown.
-            if ([self.configurator enableTransitiveRules]) {
-              cd.decision = SNTEventStateAllowTransitive;
-              return cd;
-            } else {
-              rule.state = SNTRuleStateUnknown;
-            }
-          default: break;
-        }
-        break;
-      case SNTRuleTypeSigningID:
-        switch (rule.state) {
-          case SNTRuleStateAllow: cd.decision = SNTEventStateAllowSigningID; return cd;
-          case SNTRuleStateAllowCompiler:
-            // If transitive rules are enabled, then SNTRuleStateAllowListCompiler rules
-            // become SNTEventStateAllowCompiler decisions.  Otherwise we treat the rule as if
-            // it were SNTRuleStateAllowSigningID.
-            if ([self.configurator enableTransitiveRules]) {
-              cd.decision = SNTEventStateAllowCompiler;
-            } else {
-              cd.decision = SNTEventStateAllowSigningID;
-            }
-            return cd;
-          case SNTRuleStateSilentBlock:
-            cd.silentBlock = YES;
-            // intentional fallthrough
-          case SNTRuleStateBlock:
-            cd.customMsg = rule.customMsg;
-            cd.customURL = rule.customURL;
-            cd.decision = SNTEventStateBlockSigningID;
-            return cd;
-          default: break;
-        }
-        break;
-      case SNTRuleTypeCertificate:
-        switch (rule.state) {
-          case SNTRuleStateAllow: cd.decision = SNTEventStateAllowCertificate; return cd;
-          case SNTRuleStateSilentBlock:
-            cd.silentBlock = YES;
-            // intentional fallthrough
-          case SNTRuleStateBlock:
-            cd.customMsg = rule.customMsg;
-            cd.customURL = rule.customURL;
-            cd.decision = SNTEventStateBlockCertificate;
-            return cd;
-          default: break;
-        }
-        break;
-      case SNTRuleTypeTeamID:
-        switch (rule.state) {
-          case SNTRuleStateAllow: cd.decision = SNTEventStateAllowTeamID; return cd;
-          case SNTRuleStateSilentBlock:
-            cd.silentBlock = YES;
-            // intentional fallthrough
-          case SNTRuleStateBlock:
-            cd.customMsg = rule.customMsg;
-            cd.customURL = rule.customURL;
-            cd.decision = SNTEventStateBlockTeamID;
-            return cd;
-          default: break;
-        }
-        break;
-
-      default: break;
+    // If we have a rule match we don't need to process any further.
+    if ([self decision:cd
+                      forRule:rule
+          withTransitiveRules:self.configurator.enableTransitiveRules]) {
+      return cd;
     }
   }