osquery-1/osquery/database/db_handle.h

// Copyright 2004-present Facebook. All Rights Reserved.

#ifndef OSQUERY_DATABASE_DB_HANDLE_H
#define OSQUERY_DATABASE_DB_HANDLE_H

#include <memory>
#include <string>
#include <vector>

#include "gtest/gtest_prod.h"
#include <rocksdb/db.h>

#include "osquery/core/status.h"

namespace osquery { namespace db {

/////////////////////////////////////////////////////////////////////////////
// Constants
/////////////////////////////////////////////////////////////////////////////

// kDBPath represents the path of the RocksDB database on disk
extern const std::string kDBPath;

// kDomains is a const vector of required "domains", or "groups of keys"
extern const std::vector<std::string> kDomains;

// kQueries is the "domain" where query data, the results of queries, etc is
// stored.
extern const std::string kQueries;

// kConfigurations is the "domain" where  certain global configurations are
// stored
extern const std::string kConfigurations;

/////////////////////////////////////////////////////////////////////////////
// DBHandle is a RAII singleton around RocksDB database handles.
/////////////////////////////////////////////////////////////////////////////
class DBHandle {
public:
  // DBHandle's destructor takes care of deallocating all previously allocated
  // resources
  ~DBHandle();

  // getInstance returns a singleton instance of DBHandle.
  static std::shared_ptr<DBHandle> getInstance();

  /////////////////////////////////////////////////////////////////////////////
  // getters and setters
  /////////////////////////////////////////////////////////////////////////////

  // getStatus() returns the status_ property
  osquery::core::Status getStatus();

  // getDB() is a helper that's used to get access to db_
  rocksdb::DB* getDB();

  /////////////////////////////////////////////////////////////////////////////
  // Locking methods
  /////////////////////////////////////////////////////////////////////////////

  // if a set of operations needs to be atomic, use startTransaction() before
  // your operations and endTransaction() after your operations.
  void startTransaction();
  void endTransaction();

  /////////////////////////////////////////////////////////////////////////////
  // Data manipulation methods
  /////////////////////////////////////////////////////////////////////////////

  // Get a "key" from "domain" and store it's content in "value"
  osquery::core::Status Get(
    const std::string& domain,
    const std::string& key,
    std::string& value
  );

  // Set "key" to "value" in "domain"
  osquery::core::Status Put(
    const std::string& domain,
    const std::string& key,
    const std::string& value
  );

  // Delete "key" and it's corresponding value from "domain"
  osquery::core::Status Delete(
    const std::string& domain,
    const std::string& key
  );

  // List all keys in "domain" and store the results in "results"
  osquery::core::Status Scan(
    const std::string& domain,
    std::vector<std::string>& results
  );

private:
  /////////////////////////////////////////////////////////////////////////////
  // Private methods
  /////////////////////////////////////////////////////////////////////////////

  // DBHandle's constructor takes care of properly connecting to RocksDB and
  // ensuring that all necessary column families are created. The resulting
  // database handle can then be accessed via getDB() and the success of the
  // connection can be determined by inspecting the resulting status code via
  // getStatus()
  DBHandle();
  DBHandle(std::string path, bool in_memory);

  // the private getInstance methods exist to expose a bit more of RocksDB's
  // functionality to DBHandle for use during unit tests
  static std::shared_ptr<DBHandle> getInstanceAtPath(const std::string& path);
  static std::shared_ptr<DBHandle> getInstanceInMemory();
  static std::shared_ptr<DBHandle> getInstance(
      const std::string& path, bool in_memory);

  /////////////////////////////////////////////////////////////////////////////
  // private getters and setters
  /////////////////////////////////////////////////////////////////////////////

  // getHandleForColumnFamily is a private helper around accessing the column
  // family handle for a specific column family, based on it's name
  rocksdb::ColumnFamilyHandle* getHandleForColumnFamily(const std::string& cf);

private:
  /////////////////////////////////////////////////////////////////////////////
  // Private members
  /////////////////////////////////////////////////////////////////////////////

  // db_ is the database handle
  rocksdb::DB* db_;

  // status_ is the status code that is generated during the attempt to connect
  // to RocksDB
  rocksdb::Status status_;

  // column_families_ is a vector of column family descriptors which are used
  // to connect to RocksDB
  std::vector<rocksdb::ColumnFamilyDescriptor> column_families_;

  // handles is a vector of pointers to column family handles
  std::vector<rocksdb::ColumnFamilyHandle*> handles_;

  // options_ contains the RocksDB database connection options that are used to
  // connect to RocksDB
  rocksdb::Options options_;
private:
  /////////////////////////////////////////////////////////////////////////////
  // Unit tests which can access private members
  /////////////////////////////////////////////////////////////////////////////

  FRIEND_TEST(DBHandleTests, test_create_new_database_on_disk);
  FRIEND_TEST(DBHandleTests, test_singleton_on_disk);
  FRIEND_TEST(DBHandleTests, test_get_instance_in_memory);
  FRIEND_TEST(DBHandleTests, test_is_query_name_in_database);
  FRIEND_TEST(DBHandleTests, test_get_handle_for_column_family);
  FRIEND_TEST(DBHandleTests, test_get_stored_query_names);
  FRIEND_TEST(DBHandleTests, test_get_executions);
  FRIEND_TEST(DBHandleTests, test_get);
  FRIEND_TEST(DBHandleTests, test_put);
  FRIEND_TEST(DBHandleTests, test_delete);
  FRIEND_TEST(DBHandleTests, test_scan);

  FRIEND_TEST(QueryTests, test_add_and_get_current_results);
  FRIEND_TEST(QueryTests, test_is_query_name_in_database);
  FRIEND_TEST(QueryTests, test_get_stored_query_names);
  FRIEND_TEST(QueryTests, test_get_executions);
  FRIEND_TEST(QueryTests, test_get_current_results);
  FRIEND_TEST(QueryTests, test_get_historical_query_results);
  FRIEND_TEST(QueryTests, test_query_name_not_found_in_db);
};

}}

#endif