Merge a764a6c into e7c32ae

google · May 16, 2023 · 1150272 · 1150272
2 parents e7c32ae + a764a6c
commit 1150272
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 7 deletions.
diff --git a/Source/santad/DataLayer/SNTDatabaseTable.m b/Source/santad/DataLayer/SNTDatabaseTable.m
@@ -39,9 +39,9 @@ - (instancetype)initWithDatabaseQueue:(FMDatabaseQueue *)db {
           bail = YES;
           return;
         }
-        [db close];
-        [[NSFileManager defaultManager] removeItemAtPath:[db databasePath] error:NULL];
-        [db open];
+        [self closeDeleteReopenDatabase:db];
+      } else if ([db userVersion] > [self currentSupportedVersion]) {
+        [self closeDeleteReopenDatabase:db];
       }
     }];
 
@@ -58,11 +58,22 @@ - (instancetype)init {
   return nil;
 }
 
+- (void)closeDeleteReopenDatabase:(FMDatabase *)db {
+  [db close];
+  [[NSFileManager defaultManager] removeItemAtPath:[db databasePath] error:NULL];
+  [db open];
+}
+
 - (uint32_t)initializeDatabase:(FMDatabase *)db fromVersion:(uint32_t)version {
   [self doesNotRecognizeSelector:_cmd];
   return 0;
 }
 
+- (uint32_t)currentSupportedVersion {
+  [self doesNotRecognizeSelector:_cmd];
+  return 0;
+}
+
 /// Called at the end of initialization to ensure the table in the
 /// database exists and uses the latest schema.
 - (void)updateTableSchema {

diff --git a/Source/santad/DataLayer/SNTEventTable.m b/Source/santad/DataLayer/SNTEventTable.m
@@ -18,8 +18,14 @@
 
 #import "Source/common/SNTStoredEvent.h"
 
+static const uint32_t kEventTableCurrentVersion = 3;
+
 @implementation SNTEventTable
 
+- (uint32_t)currentSupportedVersion {
+  return kEventTableCurrentVersion;
+}
+
 - (uint32_t)initializeDatabase:(FMDatabase *)db fromVersion:(uint32_t)version {
   int newVersion = 0;
 

diff --git a/Source/santad/DataLayer/SNTRuleTable.m b/Source/santad/DataLayer/SNTRuleTable.m
@@ -25,6 +25,8 @@
 #import "Source/common/SNTLogging.h"
 #import "Source/common/SNTRule.h"
 
+static const uint32_t kRuleTableCurrentVersion = 4;
+
 // TODO(nguyenphillip): this should be configurable.
 // How many rules must be in database before we start trying to remove transitive rules.
 static const NSUInteger kTransitiveRuleCullingThreshold = 500000;
@@ -173,6 +175,10 @@ - (void)setupSystemCriticalBinaries {
   self.criticalSystemBinaries = bins;
 }
 
+- (uint32_t)currentSupportedVersion {
+  return kRuleTableCurrentVersion;
+}
+
 - (uint32_t)initializeDatabase:(FMDatabase *)db fromVersion:(uint32_t)version {
   // Lock this database from other processes
   [[db executeQuery:@"PRAGMA locking_mode = EXCLUSIVE;"] close];

diff --git a/Source/santad/SNTDatabaseController.m b/Source/santad/SNTDatabaseController.m
@@ -39,8 +39,6 @@ + (SNTEventTable *)eventTable {
     NSString *fullPath =
       [[SNTDatabaseController databasePath] stringByAppendingPathComponent:kEventsDatabaseName];
     FMDatabaseQueue *dbq = [[FMDatabaseQueue alloc] initWithPath:fullPath];
-    chown([fullPath UTF8String], 0, 0);
-    chmod([fullPath UTF8String], 0600);
 
 #ifndef DEBUG
     [dbq inDatabase:^(FMDatabase *db) {
@@ -49,6 +47,9 @@ + (SNTEventTable *)eventTable {
 #endif
 
     eventDatabase = [[SNTEventTable alloc] initWithDatabaseQueue:dbq];
+
+    chown([fullPath UTF8String], 0, 0);
+    chmod([fullPath UTF8String], 0600);
   });
 
   return eventDatabase;
@@ -62,8 +63,6 @@ + (SNTRuleTable *)ruleTable {
     NSString *fullPath =
       [[SNTDatabaseController databasePath] stringByAppendingPathComponent:kRulesDatabaseName];
     FMDatabaseQueue *dbq = [[FMDatabaseQueue alloc] initWithPath:fullPath];
-    chown([fullPath UTF8String], 0, 0);
-    chmod([fullPath UTF8String], 0600);
 
 #ifndef DEBUG
     [dbq inDatabase:^(FMDatabase *db) {
@@ -72,6 +71,9 @@ + (SNTRuleTable *)ruleTable {
 #endif
 
     ruleDatabase = [[SNTRuleTable alloc] initWithDatabaseQueue:dbq];
+
+    chown([fullPath UTF8String], 0, 0);
+    chmod([fullPath UTF8String], 0600);
   });
   return ruleDatabase;
 }