Introduce scheduler reload feature (#2917)

include/osquery/database.h

@@ -589,6 +589,9 @@ Status scanDatabaseKeys(const std::string& domain,
                         const std::string& prefix,
                         size_t max = 0);
 
+/// Allow callers to reload or reset the database plugin.
+void resetDatabase();
+
 /// Allow callers to scan each column family and print each value.
 void dumpDatabase();
 }

include/osquery/logger.h

@@ -20,6 +20,8 @@
 
 #include <glog/logging.h>
 
+#include <boost/noncopyable.hpp>
+
 #include <osquery/database.h>
 #include <osquery/flags.h>
 #include <osquery/registry.h>
@@ -330,7 +332,7 @@ Status logSnapshotQuery(const QueryLogItem& item);
  * The logger forwarding state is restored and unlocked as soon as the object
  * of this class goes out of scope.
  */
-class LoggerForwardingDisabler {
+class LoggerForwardingDisabler : private boost::noncopyable {
  public:
   LoggerForwardingDisabler();
   ~LoggerForwardingDisabler();

osquery/core/posix/process_ops.cpp

@@ -35,7 +35,7 @@ int platformGetUid() {
 
 bool isLauncherProcessDead(PlatformProcess& launcher) {
   if (!launcher.isValid()) {
-    return false;
+    return true;
   }
 
   return (::getppid() != launcher.nativeHandle());

osquery/database/database.cpp

@@ -33,6 +33,7 @@ CLI_FLAG(string,
 FLAG_ALIAS(std::string, db_path, database_path);
 
 FLAG(bool, disable_database, false, "Disable the persistent RocksDB storage");
+
 DECLARE_bool(decorations_top_level);
 
 #if defined(SKIP_ROCKSDB)
@@ -54,6 +55,14 @@ bool DatabasePlugin::kDBHandleOptionAllowOpen(false);
 bool DatabasePlugin::kDBHandleOptionRequireWrite(false);
 std::atomic<bool> DatabasePlugin::kCheckingDB(false);
 
+/**
+ * @brief A reader/writer mutex protecting database resets.
+ *
+ * A write is locked while using reset flows. A read is locked when calling
+ * database plugin APIs.
+ */
+Mutex kDatabaseReset;
+
 Status serializeRow(const Row& r, pt::ptree& tree) {
   try {
     for (auto& i : r) {
@@ -447,6 +456,7 @@ void DatabasePlugin::shutdown() {
 }
 
 Status DatabasePlugin::reset() {
+  // Keep this simple, scope the critical section to the broader methods.
   tearDown();
   return setUp();
 }
@@ -505,6 +515,8 @@ Status DatabasePlugin::call(const PluginRequest& request,
       response.push_back({{"k", k}});
     }
     return status;
+  } else if (request.at("action") == "reset") {
+    return this->reset();
   }
 
   return Status(1, "Unknown database plugin action");
@@ -523,6 +535,8 @@ static inline std::shared_ptr<DatabasePlugin> getDatabasePlugin() {
 Status getDatabaseValue(const std::string& domain,
                         const std::string& key,
                         std::string& value) {
+  ReadLock lock(kDatabaseReset);
+
   if (RegistryFactory::get().external()) {
     // External registries (extensions) do not have databases active.
     // It is not possible to use an extension-based database.
@@ -546,6 +560,8 @@ Status getDatabaseValue(const std::string& domain,
 Status setDatabaseValue(const std::string& domain,
                         const std::string& key,
                         const std::string& value) {
+  ReadLock lock(kDatabaseReset);
+
   if (RegistryFactory::get().external()) {
     // External registries (extensions) do not have databases active.
     // It is not possible to use an extension-based database.
@@ -559,6 +575,8 @@ Status setDatabaseValue(const std::string& domain,
 }
 
 Status deleteDatabaseValue(const std::string& domain, const std::string& key) {
+  ReadLock lock(kDatabaseReset);
+
   if (RegistryFactory::get().external()) {
     // External registries (extensions) do not have databases active.
     // It is not possible to use an extension-based database.
@@ -582,6 +600,8 @@ Status scanDatabaseKeys(const std::string& domain,
                         std::vector<std::string>& keys,
                         const std::string& prefix,
                         size_t max) {
+  ReadLock lock(kDatabaseReset);
+
   if (RegistryFactory::get().external()) {
     // External registries (extensions) do not have databases active.
     // It is not possible to use an extension-based database.
@@ -604,6 +624,18 @@ Status scanDatabaseKeys(const std::string& domain,
   }
 }
 
+void resetDatabase() {
+  WriteLock lock(kDatabaseReset);
+
+  // Prevent RocksDB reentrancy by logger plugins during plugin setup.
+  LoggerForwardingDisabler disable_logging;
+  PluginRequest request = {{"action", "reset"}};
+  if (!Registry::call("database", request)) {
+    LOG(WARNING) << "Unable to reset database plugin: "
+                 << Registry::get().getActive("database");
+  }
+}
+
 void dumpDatabase() {
   for (const auto& domain : kDomains) {
     std::vector<std::string> keys;

osquery/database/plugins/rocksdb.cpp

@@ -130,10 +130,10 @@ void GlogRocksDBLogger::Logv(const char* format, va_list ap) {
 
   // There is a spurious warning on first open.
   if (log_line.find("Error when reading") == std::string::npos) {
-    // Rocksdb calls are non-reentrant. Since this callback is made in the
-    // context of a rocksdb api call, turn log forwarding off to prevent the
-    // logger from trying to make a call back into rocksdb and causing a
-    // deadlock
+    // RocksDB calls are non-reentrant. Since this callback is made in the
+    // context of a RocksDB API call, turn log forwarding off to prevent the
+    // logger from trying to make a call back into RocksDB and causing a
+    // deadlock.
     LoggerForwardingDisabler forwarding_disabler;
     LOG(INFO) << "RocksDB: " << log_line;
   }
@@ -141,7 +141,7 @@ void GlogRocksDBLogger::Logv(const char* format, va_list ap) {
 
 Status RocksDBDatabasePlugin::setUp() {
   if (!kDBHandleOptionAllowOpen) {
-    LOG(WARNING) << RLOG(1629) << "Not allowed to create DBHandle instance";
+    LOG(WARNING) << RLOG(1629) << "Not allowed to set up database plugin";
   }
 
   if (!initialized_) {

osquery/database/plugins/tests/rocksdb_tests.cpp

@@ -16,7 +16,9 @@ namespace osquery {
 
 class RocksDBDatabasePluginTests : public DatabasePluginTests {
  protected:
-  std::string name() override { return "rocksdb"; }
+  std::string name() override {
+    return "rocksdb";
+  }
 };
 
 // Define the default set of database plugin operation tests.

osquery/database/plugins/tests/sqlite_tests.cpp

@@ -14,7 +14,9 @@ namespace osquery {
 
 class SQLiteDatabasePluginTests : public DatabasePluginTests {
  protected:
-  std::string name() override { return "sqlite"; }
+  std::string name() override {
+    return "sqlite";
+  }
 };
 
 // Define the default set of database plugin operation tests.

osquery/database/tests/plugin_tests.cpp

@@ -38,6 +38,19 @@ void DatabasePluginTests::testPluginCheck() {
   EXPECT_TRUE(db_plugin->reset());
 }
 
+void DatabasePluginTests::testReset() {
+  RegistryFactory::get().setActive("database", getName());
+  setDatabaseValue(kLogs, "reset", "1");
+  resetDatabase();
+
+  if ("ephemeral" != getName()) {
+    // The ephemeral plugin is special and does not persist after reset.
+    std::string value;
+    EXPECT_TRUE(getDatabaseValue(kLogs, "reset", value));
+    EXPECT_EQ(value, "1");
+  }
+}
+
 void DatabasePluginTests::testPut() {
   auto s = getPlugin()->put(kQueries, "test_put", "bar");
   EXPECT_TRUE(s.ok());

osquery/database/tests/plugin_tests.h

@@ -24,6 +24,9 @@
   TEST_F(n, test_plugin_check) {                                               \
     testPluginCheck();                                                         \
   }                                                                            \
+  TEST_F(n, test_reset) {                                                      \
+    testReset();                                                               \
+  }                                                                            \
   TEST_F(n, test_put) {                                                        \
     testPut();                                                                 \
   }                                                                            \
@@ -97,6 +100,7 @@ class DatabasePluginTests : public testing::Test {
 
  protected:
   void testPluginCheck();
+  void testReset();
   void testPut();
   void testGet();
   void testDelete();

osquery/dispatcher/scheduler.cpp

@@ -21,12 +21,18 @@
 #include "osquery/core/process.h"
 #include "osquery/database/query.h"
 #include "osquery/dispatcher/scheduler.h"
+#include "osquery/sql/sqlite_util.h"
 
 namespace osquery {
 
 FLAG(bool, enable_monitor, true, "Enable the schedule monitor");
 
-FLAG(uint64, schedule_timeout, 0, "Limit the schedule, 0 for no limit")
+FLAG(uint64, schedule_timeout, 0, "Limit the schedule, 0 for no limit");
+
+FLAG(uint64,
+     schedule_reload,
+     7200,
+     "Interval in seconds to reload database arenas");
 
 /// Used to bypass (optimize-out) the set-differential of query results.
 DECLARE_bool(events_optimize);
@@ -149,9 +155,14 @@ void SchedulerRunner::start() {
           }
         }));
     // Configuration decorators run on 60 second intervals only.
-    if (i % 60 == 0) {
+    if ((i % 60) == 0) {
       runDecorators(DECORATE_INTERVAL, i);
     }
+    if (FLAGS_schedule_reload > 0 && (i % FLAGS_schedule_reload) == 0) {
+      SQLiteDBManager::resetPrimary();
+      resetDatabase();
+    }
+
     // Put the thread into an interruptible sleep without a config instance.
     pauseMilli(interval_ * 1000);
     if (interrupted()) {

osquery/dispatcher/tests/scheduler_tests.cpp

@@ -20,6 +20,7 @@
 namespace osquery {
 
 DECLARE_bool(disable_logging);
+DECLARE_uint64(schedule_reload);
 
 class SchedulerTests : public testing::Test {
   void SetUp() override {
@@ -171,4 +172,17 @@ TEST_F(SchedulerTests, test_scheduler) {
   TablePlugin::kCacheStep = backup_step;
   TablePlugin::kCacheInterval = backup_interval;
 }
+
+TEST_F(SchedulerTests, test_scheduler_reload) {
+  std::string config =
+      "{\"schedule\":{\"1\":{"
+      "\"query\":\"select * from processes\", \"interval\":1}}}";
+  auto backup_reload = FLAGS_schedule_reload;
+
+  // Start the scheduler;
+  auto expire = static_cast<unsigned long int>(getUnixTime() + 1);
+  FLAGS_schedule_reload = 1;
+  SchedulerRunner runner(expire, 1);
+  FLAGS_schedule_reload = backup_reload;
+}
 }

osquery/logger/logger.cpp

@@ -224,7 +224,7 @@ class BufferedLogSink : public google::LogSink, private boost::noncopyable {
 };
 
 /// Scoped helper to perform logging actions without races.
-class LoggerDisabler {
+class LoggerDisabler : private boost::noncopyable {
  public:
   LoggerDisabler()
       : stderr_status_(FLAGS_logtostderr),

osquery/sql/sqlite_util.cpp

@@ -251,7 +251,7 @@ void SQLiteDBInstance::clearAffectedTables() {
 }
 
 SQLiteDBInstance::~SQLiteDBInstance() {
-  if (!isPrimary()) {
+  if (!isPrimary() && db_ != nullptr) {
     sqlite3_close(db_);
   } else {
     db_ = nullptr;
@@ -268,6 +268,16 @@ bool SQLiteDBManager::isDisabled(const std::string& table_name) {
   return (element != instance().disabled_tables_.end());
 }
 
+void SQLiteDBManager::resetPrimary() {
+  auto& self = instance();
+  WriteLock create_lock(self.create_mutex_);
+  WriteLock connection_lock(self.mutex_);
+
+  self.connection_.reset();
+  sqlite3_close(self.db_);
+  self.db_ = nullptr;
+}
+
 void SQLiteDBManager::setDisabledTables(const std::string& list) {
   const auto& tables = split(list, ",");
   disabled_tables_ =

osquery/sql/sqlite_util.h

@@ -142,6 +142,14 @@ class SQLiteDBManager : private boost::noncopyable {
   /// See `get` but always return a transient DB connection (for testing).
   static SQLiteDBInstanceRef getUnique();
 
+  /**
+   * @brief Reset the primary database connection.
+   *
+   * Over time it may be helpful to remove SQLite's arena.
+   * We can periodically close and re-initialize and connect virtual tables.
+   */
+  static void resetPrimary();
+
   /**
    * @brief Check if `table_name` is disabled.
    *

osquery/sql/tests/sqlite_util_tests.cpp

@@ -64,6 +64,16 @@ TEST_F(SQLiteUtilTests, test_sqlite_instance) {
   EXPECT_EQ(internal_db, SQLiteDBManager::get()->db());
 }
 
+TEST_F(SQLiteUtilTests, test_reset) {
+  auto internal_db = SQLiteDBManager::get()->db();
+  SQLiteDBManager::resetPrimary();
+  auto new_internal_db = SQLiteDBManager::get()->db();
+
+  // Assume the internal (primary) database we reset and recreated.
+  EXPECT_NE(nullptr, new_internal_db);
+  EXPECT_NE(internal_db, new_internal_db);
+}
+
 TEST_F(SQLiteUtilTests, test_direct_query_execution) {
   auto dbc = getTestDBC();
   QueryData results;