[#1527] Add a --profile option to the shell, replace 'run'

CMakeLists.txt

@@ -44,6 +44,8 @@ add_compile_options(
   -fPIC
   -fpic
   -pipe
+  -fdata-sections
+  -ffunction-sections
 )
 
 # osquery additional compiler flags added by CMake.

osquery/CMakeLists.txt

@@ -47,11 +47,16 @@ endforeach()
 
 # The platform-specific SDK + core libraries.
 if(APPLE)
+  ADD_OSQUERY_LINK_CORE("-Wl,-dead_strip")
   ADD_OSQUERY_LINK_CORE("-mmacosx-version-min=${OSX_VERSION_MIN}")
   ADD_OSQUERY_LINK_CORE("boost_thread-mt")
   ADD_OSQUERY_LINK_CORE("lz4")
 else()
   ADD_OSQUERY_LINK_CORE("-Wl,-zrelro -Wl,-znow")
+  if(NOT ${OSQUERY_BUILD_DISTRO} STREQUAL "rhel6" AND
+     NOT ${OSQUERY_BUILD_DISTRO} STREQUAL "centos6")
+    ADD_OSQUERY_LINK_CORE("-Wl,--gc-sections")
+  endif()
   ADD_OSQUERY_LINK_CORE("boost_thread")
   ADD_OSQUERY_LINK_CORE("librt.so")
 endif()
@@ -248,13 +253,6 @@ if(NOT DEFINED ENV{SKIP_TESTS})
     target_link_libraries(osquery_tables_tests gtest libosquery_testing)
     SET_OSQUERY_COMPILE(osquery_tables_tests "${CXX_COMPILE_FLAGS} -DGTEST_HAS_TR1_TUPLE=0")
     add_test(osquery_tables_tests osquery_tables_tests)
-
-    # osquery table run profiler built outside of SDK.
-    add_executable(run main/run.cpp)
-    TARGET_OSQUERY_LINK_WHOLE(run libosquery)
-    TARGET_OSQUERY_LINK_WHOLE(run libosquery_additional)
-    SET_OSQUERY_COMPILE(run "${CXX_COMPILE_FLAGS}")
-    set_target_properties(run PROPERTIES OUTPUT_NAME run)
   endif()
 
   # Build the example extension with the SDK.

osquery/main/run.cpp

@@ -1,74 +0,0 @@
-/*
- *  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 <errno.h>
-
-#include <osquery/core.h>
-#include <osquery/events.h>
-#include <osquery/flags.h>
-#include <osquery/logger.h>
-
-#include "osquery/database/db_handle.h"
-#include "osquery/sql/sqlite_util.h"
-
-DEFINE_string(query, "", "query to execute");
-DEFINE_int32(iterations, 1, "times to run the query in question");
-DEFINE_int32(delay, 0, "delay before and after the query");
-
-namespace osquery {
-
-DECLARE_bool(disable_events);
-DECLARE_bool(registry_exceptions);
-}
-
-int main(int argc, char* argv[]) {
-  // Only log to stderr
-  FLAGS_logtostderr = true;
-
-  // Let gflags parse the non-help options/flags.
-  GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, false);
-  google::InitGoogleLogging(argv[0]);
-
-  if (FLAGS_query == "") {
-    fprintf(stderr, "Usage: %s --query=\"query\"\n", argv[0]);
-    return 1;
-  }
-
-  osquery::DBHandle::setAllowOpen(true);
-  osquery::FLAGS_database_path = "/dev/null";
-  osquery::Registry::setUp();
-  osquery::FLAGS_disable_events = true;
-  osquery::FLAGS_registry_exceptions = true;
-  osquery::attachEvents();
-
-  if (FLAGS_delay != 0) {
-    ::sleep(FLAGS_delay);
-  }
-
-  osquery::QueryData results;
-  osquery::Status status;
-  for (int i = 0; i < FLAGS_iterations; ++i) {
-    auto dbc = osquery::SQLiteDBManager::get();
-    status = osquery::queryInternal(FLAGS_query, results, dbc.db());
-    if (!status.ok()) {
-      fprintf(stderr, "Query failed: %d\n", status.getCode());
-      break;
-    }
-  }
-
-  if (FLAGS_delay != 0) {
-    ::sleep(FLAGS_delay);
-  }
-
-  // Instead of calling "shutdownOsquery" force the EF to join its threads.
-  GFLAGS_NAMESPACE::ShutDownCommandLineFlags();
-
-  return status.getCode();
-}

osquery/main/shell.cpp

@@ -12,16 +12,66 @@
 
 #include <osquery/core.h>
 #include <osquery/extensions.h>
+#include <osquery/flags.h>
 
 #include "osquery/core/watcher.h"
 #include "osquery/devtools/devtools.h"
+#include "osquery/sql/sqlite_util.h"
+
+namespace osquery {
+
+SHELL_FLAG(int32,
+           profile,
+           0,
+           "Enable profile mode when non-0, set number of iterations");
+HIDDEN_FLAG(int32,
+            profile_delay,
+            0,
+            "Sleep a number of seconds before and after the profiling");
+}
+
+int profile(int argc, char *argv[]) {
+  std::string query;
+  if (!isatty(fileno(stdin))) {
+    std::getline(std::cin, query);
+  } else if (argc < 2) {
+    // No query input provided via stdin or as a positional argument.
+    fprintf(stderr, "No query provided via stdin or args to profile...\n");
+    return 2;
+  } else {
+    query = std::string(argv[1]);
+  }
+
+  if (osquery::FLAGS_profile_delay > 0) {
+    ::sleep(osquery::FLAGS_profile_delay);
+  }
+
+  auto dbc = osquery::SQLiteDBManager::get();
+  for (size_t i = 0; i < static_cast<size_t>(osquery::FLAGS_profile); ++i) {
+    osquery::QueryData results;
+    auto status = osquery::queryInternal(query, results, dbc.db());
+    if (!status) {
+      fprintf(stderr,
+              "Query failed (%d): %s\n",
+              status.getCode(),
+              status.what().c_str());
+      return status.getCode();
+    }
+  }
+
+  if (osquery::FLAGS_profile_delay > 0) {
+    ::sleep(osquery::FLAGS_profile_delay);
+  }
+
+  return 0;
+}
 
 int main(int argc, char *argv[]) {
   // Parse/apply flags, start registry, load logger/config plugins.
   osquery::Initializer runner(argc, argv, osquery::OSQUERY_TOOL_SHELL);
   if (argc > 1 || !isatty(fileno(stdin)) || osquery::FLAGS_A.size() > 0 ||
-      osquery::FLAGS_L) {
-    // A query was set as a positional argument for via stdin.
+      osquery::FLAGS_L || osquery::FLAGS_profile > 0) {
+    // A query was set as a positional argument, via stdin, or profiling is on.
     osquery::FLAGS_disable_events = true;
     // The shell may have loaded table extensions, if not, disable the manager.
     if (!osquery::Watcher::hasManagedExtensions()) {
@@ -29,10 +79,15 @@ int main(int argc, char *argv[]) {
     }
   }
 
-  runner.start();
+  int retcode = 0;
+  if (osquery::FLAGS_profile <= 0) {
+    runner.start();
 
-  // Virtual tables will be attached to the shell's in-memory SQLite DB.
-  int retcode = osquery::launchIntoShell(argc, argv);
+    // Virtual tables will be attached to the shell's in-memory SQLite DB.
+    retcode = osquery::launchIntoShell(argc, argv);
+  } else {
+    retcode = profile(argc, argv);
+  }
 
   // Finally shutdown.
   runner.shutdown();

tools/profile.py

@@ -65,7 +65,7 @@ def check_leaks_linux(shell, query, count=1, supp_file=None):
 def check_leaks_darwin(shell, query, count=1):
     # Run the shell with a --delay flag such that leaks can attach before exit.
     proc = subprocess.Popen(
-        [shell, "--query", query, "--iterations", str(count), "--delay", "1"],
+        [shell, "--profile", str(count), "--profile_delay", "1", query],
         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     leak_checks = None
     while proc.poll() is None:
@@ -119,20 +119,15 @@ def profile_leaks(shell, queries, count=1, rounds=1, supp_file=None):
 
 
 def run_query(shell, query, timeout=0, count=1):
-    """Execute the osquery run testing wrapper with a setup/teardown delay."""
+    """Execute the osqueryi shell in profile mode with a setup/teardown delay."""
     start_time = time.time()
-    proc = subprocess.Popen(
-        [shell, "--query", query, "--iterations", str(count),
-            "--delay", "1"],
-        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     return utils.profile_cmd([
         shell,
-        "--query",
-        query,
-        "--iterations",
+        "--profile",
         str(count),
-        "--delay",
-        "1"
+        "--profile_delay",
+        "1",
+        query
     ])
 
 
@@ -272,9 +267,9 @@ if __name__ == "__main__":
         help="Run the profile for N rounds and use the average."
     )
     group.add_argument(
-        "--shell", metavar="PATH", default="./build/%s/osquery/run" % (
+        "--shell", metavar="PATH", default="./build/%s/osquery/osqueryi" % (
             utils.platform()),
-        help="Path to osquery run wrapper (./build/<sys>/osquery/run)."
+        help="Path to osqueryi shell (./build/<sys>/osquery/osqueryi)."
     )
 
     group = parser.add_argument_group("Performance Options:")

tools/tests/test_example_queries.py

@@ -24,12 +24,15 @@ import test_base
 import utils
 
 class ExampleQueryTests(test_base.QueryTester):
+    @test_base.flaky
     def test_cross_platform_queries(self):
         self._execute_set(PLATFORM_EXAMPLES["specs"])
 
+    @test_base.flaky
     def test_platform_specific_queries(self):
         self._execute_set(PLATFORM_EXAMPLES[utils.platform()])
 
+    @test_base.flaky
     def test_utility_queries(self):
         self._execute_set(PLATFORM_EXAMPLES["utility"])