Check extensions are registered with osquery core (#6374)

osquery/extensions/extensions.cpp

@@ -13,6 +13,7 @@
 #include <vector>
 
 #include <boost/algorithm/string/trim.hpp>
+#include <boost/optional.hpp>
 
 #include <osquery/core.h>
 #include <osquery/core/watcher.h>
@@ -97,6 +98,10 @@ EXTENSION_FLAG_ALIAS(interval, extensions_interval);
 class ExtensionWatcher : public InternalRunnable {
  public:
   virtual ~ExtensionWatcher() = default;
+  ExtensionWatcher(const std::string& path,
+                   size_t interval,
+                   bool fatal,
+                   RouteUUID uuid);
   ExtensionWatcher(const std::string& path, size_t interval, bool fatal);
 
  public:
@@ -119,6 +124,10 @@ class ExtensionWatcher : public InternalRunnable {
 
   /// If the ExtensionManager socket is closed, should the extension exit.
   bool fatal_;
+
+  /// Optional uuid used to monitor if socket is registered on with extension
+  /// core
+  boost::optional<RouteUUID> uuid_;
 };
 
 class ExtensionManagerWatcher : public ExtensionWatcher {
@@ -181,6 +190,14 @@ Status extensionPathActive(const std::string& path, bool use_timeout = false) {
   }));
 }
 
+ExtensionWatcher::ExtensionWatcher(const std::string& path,
+                                   size_t interval,
+                                   bool fatal,
+                                   RouteUUID uuid)
+    : ExtensionWatcher(path, interval, fatal) {
+  uuid_ = uuid;
+}
+
 ExtensionWatcher::ExtensionWatcher(const std::string& path,
                                    size_t interval,
                                    bool fatal)
@@ -237,9 +254,31 @@ void ExtensionWatcher::watch() {
   bool core_sane = true;
   if (socketExists(path_)) {
     try {
-      ExtensionManagerClient client(path_);
-      // Ping the extension manager until it goes down.
-      status = client.ping();
+      if (uuid_) {
+        // Check we are still registered with the osquery core by getting
+        // available extensions
+        ExtensionList extensionList;
+        status = getExtensions(extensionList);
+        if (status.getCode() != (int)ExtensionCode::EXT_SUCCESS && fatal_) {
+          LOG(ERROR) << "Extension watcher failed to get extensions: "
+                     << status.getMessage();
+          exitFatal();
+        } else if (extensionList.find(uuid_.get()) == extensionList.end() &&
+                   fatal_) {
+          LOG(ERROR) << "Extension not registered with osquery core";
+          exitFatal();
+        }
+      } else {
+        // Ping the extension manager to check it's still there
+        ExtensionManagerClient client(path_);
+        status = client.ping();
+        if (status.getCode() != (int)ExtensionCode::EXT_SUCCESS && fatal_) {
+          // The core may be healthy but return a failed ping status.
+          LOG(ERROR) << "Extension watcher ping failed: "
+                     << status.getMessage();
+          exitFatal();
+        }
+      }
     } catch (const std::exception& /* e */) {
       core_sane = false;
     }
@@ -249,14 +288,9 @@ void ExtensionWatcher::watch() {
   }
 
   if (!core_sane) {
-    VLOG(1) << "Extension watcher ending: osquery core has gone away";
+    LOG(INFO) << "Extension watcher ending: osquery core has gone away";
     exitFatal(0);
   }
-
-  if (status.getCode() != (int)ExtensionCode::EXT_SUCCESS && fatal_) {
-    // The core may be healthy but return a failed ping status.
-    exitFatal();
-  }
 }
 
 void ExtensionManagerWatcher::watch() {
@@ -448,15 +482,7 @@ Status startExtension(const std::string& name,
   // When a broadcast is requested this registry should not send core plugins.
   RegistryFactory::get().setExternal();
 
-  // Latency converted to milliseconds, used as a thread interruptible.
-  auto latency = atoi(FLAGS_extensions_interval.c_str()) * 1000;
-  auto status = startExtensionWatcher(FLAGS_extensions_socket, latency, true);
-  if (!status.ok()) {
-    // If the threaded watcher fails to start, fail the extension.
-    return status;
-  }
-
-  status = startExtension(
+  auto status = startExtension(
       FLAGS_extensions_socket, name, version, min_sdk_version, kSDKVersion);
   if (!status.ok()) {
     // If the extension failed to start then the EM is most likely unavailable.
@@ -508,6 +534,16 @@ Status startExtension(const std::string& manager_path,
   // Now that the UUID is known, try to clean up stale socket paths.
   auto extension_path = getExtensionSocket(uuid, manager_path);
 
+  // Latency converted to milliseconds, used as a thread interruptible.
+  auto latency = atoi(FLAGS_extensions_interval.c_str()) * 1000;
+
+  // Register the watcher with it's uuid
+  status = startExtensionWatcher(manager_path, latency, true, uuid);
+  if (!status.ok()) {
+    // If the threaded watcher fails to start, fail the extension.
+    return status;
+  }
+
   status = socketExists(extension_path, true);
   if (!status) {
     return status;
@@ -668,7 +704,8 @@ Status callExtension(const std::string& extension_path,
 
 Status startExtensionWatcher(const std::string& manager_path,
                              size_t interval,
-                             bool fatal) {
+                             bool fatal,
+                             RouteUUID uuid) {
   // Make sure the extension manager path exists, and is writable.
   auto status = extensionPathActive(manager_path, true);
   if (!status.ok()) {
@@ -677,7 +714,7 @@ Status startExtensionWatcher(const std::string& manager_path,
 
   // Start a extension watcher, if the manager dies, so should we.
   Dispatcher::addService(
-      std::make_shared<ExtensionWatcher>(manager_path, interval, fatal));
+      std::make_shared<ExtensionWatcher>(manager_path, interval, fatal, uuid));
   return Status::success();
 }
 
@@ -748,4 +785,4 @@ Status startExtensionManager(const std::string& manager_path) {
 
   return Status::success();
 }
-} // namespace osquery
+} // namespace osquery

osquery/include/osquery/extensions.h

@@ -10,8 +10,8 @@
 
 #include <osquery/core.h>
 #include <osquery/flags.h>
-#include <osquery/registry_interface.h>
 #include <osquery/plugins/sql.h>
+#include <osquery/registry_interface.h>
 
 namespace osquery {
 
@@ -176,11 +176,12 @@ Status startExtension(const std::string& manager_path,
 /// Start an ExtensionWatcher thread.
 Status startExtensionWatcher(const std::string& manager_path,
                              size_t interval,
-                             bool fatal);
+                             bool fatal,
+                             RouteUUID uuid);
 
 /// Start an ExtensionManagerRunner thread.
 Status startExtensionManager();
 
 /// Internal startExtensionManager implementation.
 Status startExtensionManager(const std::string& manager_path);
-}
+} // namespace osquery