Plugin routable registry (via thrift)

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 <memory>
#include <string>

#include <boost/noncopyable.hpp>

#include <gtest/gtest.h>

#include <osquery/logger.h>

namespace osquery {

class RegistryTests : public testing::Test {};

template <class RegistryType>
class RegistryFactory : private boost::noncopyable {
 protected:
  typedef typename std::shared_ptr<RegistryType> RegistryTypeRef;
  typedef RegistryFactory<RegistryType> AutoRegistryType;

 public:
  static RegistryFactory& instance() {
    static RegistryFactory instance;
    return instance;
  }

  template <class Item>
  static RegistryFactory& add(const std::string& name) {
    if (instance().registries_.count(name) > 0) {
      return instance();
    }
    auto item = std::make_shared<Item>();
    item->init();
    instance().items_[name] = reinterpret_cast<RegistryTypeRef&>(item);
    return instance();
  }

  static RegistryTypeRef get(const std::string& name) {
    if (instance().items_.count(name) > 0) {
      return instance().items_[name];
    }
    return 0;
  }

  template <class Type>
  static RegistryFactory<Type>& create(const std::string& name) {
    auto& registry = RegistryFactory<Type>::instance();
    instance().registries_[name] = &reinterpret_cast<AutoRegistryType&>(registry);
    return registry;
  }

  static AutoRegistryType& registry(const std::string& name) {
    return *instance().registries_[name];
  }

  static const std::map<std::string, RegistryTypeRef>& all() {
    return instance().items_;
  }

  static size_t count() {
    return instance().items_.size();
  }

 protected:
  RegistryFactory() {}

 private:
  std::map<std::string, RegistryTypeRef> items_;
  std::map<std::string, AutoRegistryType*> registries_;
};

class Plugin {
 public:
  /// The plugin may perform some initialization, not required.
  virtual void init() {}
  /// The plugin may publish route info (other than registry type and name).
  virtual std::string routeInfo() { return ""; }
};

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

 protected:
  int some_value_;
};

class HouseCat : public CatPlugin {
 public:
  void init() {
    // Make sure the Plugin implementation's init is called.
    some_value_ = 9000;
  }
};

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

TEST_F(RegistryTests, test_core_registry) {
  /// Add a CatRegistry item (a plugin) called "house".
  CatRegistry::add<HouseCat>("house");
  EXPECT_EQ(CatRegistry::count(), 1);

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

  /// Request a plugin to call an API method.
  auto cat = CatRegistry::get("house");
  EXPECT_EQ(cat->getValue(), 9000);

  /// Now let's iterate over every registered Cat plugin.
  EXPECT_EQ(CatRegistry::all().size(), 2);
  for (const auto& cat : CatRegistry::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
/// registered 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 TestRegistry : 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 = TestRegistry::create<CatPlugin>("cat");

TEST_F(RegistryTests, test_factory) {
  /// Using the registry, and a registry type by name, we can registry a
  /// plugin HouseCat called "house" like above.
  TestRegistry::registry("cat").add<HouseCat>("house");
  /// 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(TestRegistry::registry("cat").count(), 1);

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

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