Handle remove rules.

google · Mar 21, 2024 · e1d7967 · e1d7967
1 parent 9d0834a
commit e1d7967
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 14 deletions.
diff --git a/Source/santad/DataLayer/SNTRuleTable.m b/Source/santad/DataLayer/SNTRuleTable.m
@@ -471,7 +471,7 @@ - (BOOL)addedRulesShouldFlushDecisionCache:(NSArray *)rules {
     }
   }
 
-  // Just flush if we have a lot of block rules.
+  // Just flush if we more than 1000 block rules.
   if (nonAllowRuleCount >= 1000) {
     return YES;
   }
@@ -487,24 +487,34 @@ - (BOOL)addedRulesShouldFlushDecisionCache:(NSArray *)rules {
 
   [self inTransaction:^(FMDatabase *db, BOOL *rollback) {
     for (SNTRule *rule in rules) {
-      if (rule.state != SNTRuleStateAllow) {
+      // If the rule is a block rule, check if it exists in the database.
+      if ((rule.state == SNTRuleStateBlock)) {
         if ([db longForQuery:
                   @"SELECT COUNT(*) FROM rules WHERE identifier=? AND type=? AND state=? LIMIT 1",
                   rule.identifier, @(rule.type), @(rule.state)] == 0) {
           flushDecisionCache = YES;
           break;
         }
-        // Skip to check the next rule since we only want to check allows for
-        // compiler rules.
-        continue;
       }
 
-      // Allowlist certificate rules are ignored
+      // If the rule being removed is targeting an allow rule, flush the cache.
+      if (rule.state == SNTRuleStateRemove) {
+        if ([db longForQuery:@"SELECT COUNT(*) FROM rules WHERE identifier=? AND type=? AND state "
+                             @"in (?,?) LIMIT 1",
+                             rule.identifier, @(rule.type), @(SNTRuleStateAllow),
+                             @(SNTRuleStateAllowCompiler)] > 0) {
+          flushDecisionCache = YES;
+          break;
+        }
+      }
+
+      // Skip certificate and TeamID rules as they cannot be compiler rules.
       if (rule.type == SNTRuleTypeCertificate || rule.type == SNTRuleTypeTeamID) continue;
 
-      if ([db longForQuery:
-                @"SELECT COUNT(*) FROM rules WHERE identifier=? AND type=? AND state=? LIMIT 1",
-                rule.identifier, @(SNTRuleTypeBinary), @(SNTRuleStateAllowCompiler)] > 0) {
+      if ([db longForQuery:@"SELECT COUNT(*) FROM rules WHERE identifier=? AND type IN (?, ?, ?)"
+                           @" AND state=? LIMIT 1",
+                           rule.identifier, @(SNTRuleTypeCDHash), @(SNTRuleTypeBinary),
+                           @(SNTRuleTypeSigningID), @(SNTRuleStateAllowCompiler)] > 0) {
         flushDecisionCache = YES;
         break;
       }
@@ -574,7 +584,6 @@ - (BOOL)fillError:(NSError **)error code:(SNTRuleTableError)code message:(NSStri
   *error = [NSError errorWithDomain:@"com.google.santad.ruletable" code:code userInfo:d];
   return YES;
 }
-
 #pragma mark Querying
 
 // Retrieve all rules from the Database

diff --git a/Source/santad/DataLayer/SNTRuleTableTest.m b/Source/santad/DataLayer/SNTRuleTableTest.m
@@ -308,8 +308,11 @@ - (void)testFetchCDHashRule {
 - (void)testFetchRuleOrdering {
   NSError *err;
   [self.sut addRules:@[
-    [self _exampleCertRule], [self _exampleBinaryRule], [self _exampleTeamIDRule],
-    [self _exampleSigningIDRuleIsPlatform:NO], [self _exampleCDHashRule],
+    [self _exampleCertRule],
+    [self _exampleBinaryRule],
+    [self _exampleTeamIDRule],
+    [self _exampleSigningIDRuleIsPlatform:NO],
+    [self _exampleCDHashRule],
   ]
          ruleCleanup:SNTRuleCleanupNone
                error:&err];
@@ -452,8 +455,8 @@ - (void)testAddedRulesShouldFlushDecisionCacheWithNewBlockRule {
   XCTAssertEqual(YES, [self.sut addedRulesShouldFlushDecisionCache:@[ r ]]);
 }
 
+// Ensure that a brand new block rule flushes the decision cache.
 - (void)testAddedRulesShouldFlushDecisionCacheWithOldBlockRule {
-  // Ensure that a brand new block rule flushes the decision cache.
   NSError *error;
   SNTRule *r = [self _exampleBinaryRule];
   [self.sut addRules:@[ r ] ruleCleanup:SNTRuleCleanupNone error:&error];
@@ -463,8 +466,8 @@ - (void)testAddedRulesShouldFlushDecisionCacheWithOldBlockRule {
   XCTAssertEqual(NO, [self.sut addedRulesShouldFlushDecisionCache:@[ r ]]);
 }
 
+// Ensure that a larger number of blocks flushes the decision cache.
 - (void)testAddedRulesShouldFlushDecisionCacheWithLargeNumberOfBlocks {
-  // Ensure that a brand new block rule flushes the decision cache.
   NSError *error;
   SNTRule *r = [self _exampleBinaryRule];
   [self.sut addRules:@[ r ] ruleCleanup:SNTRuleCleanupNone error:&error];
@@ -479,4 +482,35 @@ - (void)testAddedRulesShouldFlushDecisionCacheWithLargeNumberOfBlocks {
   XCTAssertEqual(YES, [self.sut addedRulesShouldFlushDecisionCache:newRules]);
 }
 
+// Ensure that an allow rule that overrides a compiler rule flushes the
+// decision cache.
+- (void)testAddedRulesShouldFlushDecisionCacheWithCompilerRule {
+  NSError *error;
+  SNTRule *r = [self _exampleBinaryRule];
+  r.type = SNTRuleTypeBinary;
+  r.state = SNTRuleStateAllowCompiler;
+  [self.sut addRules:@[ r ] ruleCleanup:SNTRuleCleanupNone error:&error];
+  XCTAssertNil(error);
+  XCTAssertEqual(self.sut.ruleCount, 1);
+  XCTAssertEqual(self.sut.binaryRuleCount, 1);
+  // make the rule an allow rule
+  r.state = SNTRuleStateAllow;
+  XCTAssertEqual(YES, [self.sut addedRulesShouldFlushDecisionCache:@[ r ]]);
+}
+
+// Ensure that an Remove rule targeting an allow rule causes a flush of the cache.
+- (void)testAddedRulesShouldFlushDecisionCacheWithRemoveRule {
+  NSError *error;
+  SNTRule *r = [self _exampleBinaryRule];
+  r.type = SNTRuleTypeBinary;
+  r.state = SNTRuleStateAllow;
+  [self.sut addRules:@[ r ] ruleCleanup:SNTRuleCleanupNone error:&error];
+  XCTAssertNil(error);
+  XCTAssertEqual(self.sut.ruleCount, 1);
+  XCTAssertEqual(self.sut.binaryRuleCount, 1);
+
+  r.state = SNTRuleStateRemove;
+  XCTAssertEqual(YES, [self.sut addedRulesShouldFlushDecisionCache:@[ r ]]);
+}
+
 @end