osquery-1/osquery/registry/registry_tests.cpp

/*
 *  Copyright (c) 2014, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant 
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */
 
#include <gtest/gtest.h>

#include <osquery/logger.h>
#include <osquery/registry.h>

namespace osquery {

class RegistryTests : public testing::Test {};

class CatPlugin : public Plugin {
 public:
  CatPlugin() : some_value_(0) {}
  virtual int getValue() { return some_value_; }
  bool sayTrue() { return true; }

 protected:
  int some_value_;
};

class HouseCat : public CatPlugin {
 public:
  Status setUp() {
    // Make sure the Plugin implementation's init is called.
    some_value_ = 9000;
    return Status(0, "OK");
  }
};

/// This is a manual registry type without a name, so we cannot broadcast
/// this registry type and it does NOT need to conform to a registry API.
class CatRegistry : public RegistryCore<CatPlugin> {};

TEST_F(RegistryTests, test_registry) {
  CatRegistry cats;

  /// Add a CatRegistry item (a plugin) called "house".
  cats.add<HouseCat>("house");
  EXPECT_EQ(cats.count(), 1);

  /// Try to add the same plugin with the same name, this is meaningless.
  cats.add<HouseCat>("house");
  /// Now add the same plugin with a different name, a new plugin instance
  /// will be created and registered.
  cats.add<HouseCat>("house2");
  EXPECT_EQ(cats.count(), 2);

  /// Request a plugin to call an API method.
  auto cat = cats.get("house");
  cats.setUp();

  EXPECT_EQ(cat->getValue(), 9000);

  /// Now let's iterate over every registered Cat plugin.
  EXPECT_EQ(cats.all().size(), 2);
  for (const auto& cat : cats.all()) {
    EXPECT_TRUE(cat.second->sayTrue());
  }
}

/// To track registry types and then broadcast them via Thrift we must define
/// a plugin API. All broadcasted registry types and the plugins of that type
/// must conform to this API.
class TestPluginAPI : public Plugin {
 public:
  virtual int getValue() = 0;
  virtual bool sayTrue() = 0;
};

/// Normally we have "Registry" that dictates the set of possible API methods
/// for all registry types. Here we use a "TestRegistry" instead.
class TestCoreRegistry : public RegistryFactory<TestPluginAPI> {};

/// We can automatically create a registry type as long as that type conforms
/// to the registry API defined in the "Registry". Here we use "TestRegistry".
/// The above "CatRegistry" was easier to understand, but using a auto
/// registry via the registry create method, we can assign a tracked name
/// and then broadcast that registry name to other plugins.
auto AutoCatRegistry = TestCoreRegistry::create<CatPlugin>("cat");

TEST_F(RegistryTests, test_auto_factory) {
  /// Using the registry, and a registry type by name, we can register a
  /// plugin HouseCat called "house" like above.
  TestCoreRegistry::registry("cat")->add<HouseCat>("auto_house");
  TestCoreRegistry::add<HouseCat>("cat", "auto_house2");
  TestCoreRegistry::registry("cat")->setUp();

  /// When acting on registries by name we can check the broadcasted
  /// registry name of other plugin processes (via Thrift) as well as
  /// internally registered plugins like HouseCat.
  EXPECT_EQ(TestCoreRegistry::registry("cat")->count(), 2);
  EXPECT_EQ(TestCoreRegistry::count("cat"), 2);

  /// And we can call an API method, since we guarantee CatPlugins conform
  /// to the "TestCoreRegistry"'s "TestPluginAPI".
  auto cat = TestCoreRegistry::get("cat", "auto_house");
  EXPECT_EQ(cat->getValue(), 9000);
}

class DogPlugin : public Plugin {
 public:
  DogPlugin() : some_value_(10000) {}
  virtual int getValue() { return some_value_; }
  bool sayTrue() { return true; }

 protected:
  int some_value_;
};

class Doge : public DogPlugin {
 public:
  Doge() { some_value_ = 100000; }
};

class BadDoge : public DogPlugin {
 public:
  Status setUp() { return Status(1, "Expect error... this is a bad dog"); }
};

auto AutoDogRegistry = TestCoreRegistry::create<DogPlugin>("dog");

TEST_F(RegistryTests, test_auto_registries) {
  TestCoreRegistry::add<Doge>("dog", "doge");
  TestCoreRegistry::registry("dog")->setUp();

  EXPECT_EQ(TestCoreRegistry::count("dog"), 1);
  EXPECT_EQ(TestCoreRegistry::get("dog", "doge")->getValue(), 100000);
}

TEST_F(RegistryTests, test_persistant_registries) {
  EXPECT_EQ(TestCoreRegistry::count("cat"), 2);
}

TEST_F(RegistryTests, test_registry_exceptions) {
  EXPECT_TRUE(TestCoreRegistry::add<Doge>("dog", "duplicate_dog").ok());
  // Bad dog will be added fine, but when setup is run, it will be removed.
  EXPECT_TRUE(TestCoreRegistry::add<BadDoge>("dog", "bad_doge").ok());
  TestCoreRegistry::registry("dog")->setUp();
  // Make sure bad dog does not exist.
  EXPECT_FALSE(TestCoreRegistry::exists("dog", "bad_doge"));
  EXPECT_EQ(TestCoreRegistry::count("dog"), 2);

  int exception_count = 0;
  try {
    TestCoreRegistry::registry("does_not_exist");
  } catch (const std::out_of_range& e) {
    exception_count++;
  }

  try {
    TestCoreRegistry::add<HouseCat>("does_not_exist", "cat");
  } catch (const std::out_of_range& e) {
    exception_count++;
  }

  EXPECT_EQ(exception_count, 2);
}

class WidgetPlugin : public Plugin {
 public:
  /// The route information will usually be provided by the plugin type.
  /// The plugin/registry item will set some structures for the plugin
  /// to parse and format. BUT a plugin/registry item can also fill this
  /// information in if the plugin type/registry type exposes routeInfo as
  /// a virtual method.
  RouteInfo routeInfo() {
    RouteInfo info;
    info["name"] = name_;
    return info;
  }

  /// Plugin types should contain generic request/response formatters and
  /// decorators.
  std::string secretPower(const PluginRequest& request) {
    if (request.count("secret_power") > 0) {
      return request.at("secret_power");
    }
    return "no_secret_power";
  }
};

class SpecialWidget : public WidgetPlugin {
 public:
  Status call(const PluginRequest& request, PluginResponse& response);
};

Status SpecialWidget::call(const PluginRequest& request,
                           PluginResponse& response) {
  response.push_back(request);
  response[0]["from"] = name_;
  response[0]["secret_power"] = secretPower(request);
  return Status(0, "OK");
}

TEST_F(RegistryTests, test_registry_api) {
  auto AutoWidgetRegistry = TestCoreRegistry::create<WidgetPlugin>("widgets");
  TestCoreRegistry::add<SpecialWidget>("widgets", "special");

  // Test route info propogation, from item to registry, to broadcast.
  auto ri = TestCoreRegistry::get("widgets", "special")->routeInfo();
  EXPECT_EQ(ri.at("name"), "special");
  auto rr = TestCoreRegistry::registry("widgets")->getRoutes();
  EXPECT_EQ(rr.size(), 1);
  EXPECT_EQ(rr.at("special").at("name"), "special");

  // Broadcast will include all registries, and all their items.
  auto broadcast_info = TestCoreRegistry::getBroadcast();
  EXPECT_EQ(broadcast_info.size(), 3);
  EXPECT_EQ(broadcast_info.at("widgets").at("special").at("name"), "special");

  PluginResponse response;
  PluginRequest request;
  auto status = TestCoreRegistry::call("widgets", "special", request, response);
  EXPECT_TRUE(status.ok());
  EXPECT_EQ(response[0].at("from"), "special");
  EXPECT_EQ(response[0].at("secret_power"), "no_secret_power");

  request["secret_power"] = "magic";
  status = TestCoreRegistry::call("widgets", "special", request, response);
  EXPECT_EQ(response[0].at("secret_power"), "magic");
}
}

int main(int argc, char* argv[]) {
  testing::InitGoogleTest(&argc, argv);
  google::InitGoogleLogging(argv[0]);
  return RUN_ALL_TESTS();
}