diff --git a/include/osquery/database.h b/include/osquery/database.h index 176a1058..010314bd 100644 --- a/include/osquery/database.h +++ b/include/osquery/database.h @@ -72,6 +72,13 @@ using RowData = std::string; */ using Row = std::map; +/** + * @brief A vector of column names associated with a query + * + * ColumnNames is a vector of the column names, in order, returned by a query. + */ +using ColumnNames = std::vector; + /** * @brief Serialize a Row into a property tree * @@ -131,6 +138,19 @@ using QueryData = std::vector; Status serializeQueryData(const QueryData& q, boost::property_tree::ptree& tree); +/** + * @brief Serialize a QueryData object into a property tree + * + * @param q the QueryData to serialize + * @param cols the TableColumn vector indicating column order + * @param tree the output property tree + * + * @return Status indicating the success or failure of the operation + */ +Status serializeQueryData(const QueryData& q, + const ColumnNames& cols, + boost::property_tree::ptree& tree); + /** * @brief Serialize a QueryData object into a JSON string * diff --git a/include/osquery/distributed.h b/include/osquery/distributed.h index dcc441eb..f8b132fc 100644 --- a/include/osquery/distributed.h +++ b/include/osquery/distributed.h @@ -83,11 +83,13 @@ struct DistributedQueryResult { DistributedQueryResult() {} DistributedQueryResult(const DistributedQueryRequest& req, const QueryData& res, + const ColumnNames& cols, const Status& s) - : request(req), results(res), status(s) {} + : request(req), results(res), columns(cols), status(s) {} DistributedQueryRequest request; QueryData results; + ColumnNames columns; Status status; }; diff --git a/include/osquery/sql.h b/include/osquery/sql.h index 9fc8ade3..e52d23e2 100644 --- a/include/osquery/sql.h +++ b/include/osquery/sql.h @@ -89,6 +89,13 @@ class SQL : private only_movable { */ const QueryData& rows() const; + /** + * @brief Column information for the query + * + * @return A ColumnNames object for the query + */ + const ColumnNames& columns(); + /** * @brief Accessor to switch off of when checking the success of a query. * @@ -146,11 +153,17 @@ class SQL : private only_movable { SQL() {} protected: + /// The internal member which holds the query string + std::string q_; + /// The internal member which holds the results of the query. QueryData results_; /// The internal member which holds the status of the query. Status status_; + + /// The internal member which holds the column names and order for the query + ColumnNames columns_; }; /** diff --git a/osquery/database/benchmarks/database_benchmarks.cpp b/osquery/database/benchmarks/database_benchmarks.cpp index b88c5fe1..ade21dae 100644 --- a/osquery/database/benchmarks/database_benchmarks.cpp +++ b/osquery/database/benchmarks/database_benchmarks.cpp @@ -33,6 +33,14 @@ QueryData getExampleQueryData(size_t x, size_t y) { return qd; } +ColumnNames getExampleColumnNames(size_t x) { + ColumnNames cn; + for (size_t i = 0; i < x; i++) { + cn.push_back("key" + std::to_string(i)); + } + return cn; +} + static void DATABASE_serialize(benchmark::State& state) { auto qd = getExampleQueryData(state.range_x(), state.range_y()); while (state.KeepRunning()) { @@ -43,6 +51,20 @@ static void DATABASE_serialize(benchmark::State& state) { BENCHMARK(DATABASE_serialize)->ArgPair(1, 1)->ArgPair(10, 10)->ArgPair(10, 100); +static void DATABASE_serialize_column_order(benchmark::State& state) { + auto qd = getExampleQueryData(state.range_x(), state.range_y()); + auto cn = getExampleColumnNames(state.range_x()); + while (state.KeepRunning()) { + boost::property_tree::ptree tree; + serializeQueryData(qd, cn, tree); + } +} + +BENCHMARK(DATABASE_serialize_column_order) + ->ArgPair(1, 1) + ->ArgPair(10, 10) + ->ArgPair(10, 100); + static void DATABASE_serialize_json(benchmark::State& state) { auto qd = getExampleQueryData(state.range_x(), state.range_y()); while (state.KeepRunning()) { diff --git a/osquery/database/database.cpp b/osquery/database/database.cpp index a4839df8..f3fd4278 100644 --- a/osquery/database/database.cpp +++ b/osquery/database/database.cpp @@ -65,6 +65,17 @@ Status serializeRow(const Row& r, pt::ptree& tree) { return Status(0, "OK"); } +Status serializeRow(const Row& r, const ColumnNames& cols, pt::ptree& tree) { + try { + for (auto& c : cols) { + tree.add(c, r.at(c)); + } + } catch (const std::exception& e) { + return Status(1, e.what()); + } + return Status(0, "OK"); +} + Status serializeRowJSON(const Row& r, std::string& json) { pt::ptree tree; auto status = serializeRow(r, tree); @@ -116,6 +127,20 @@ Status serializeQueryData(const QueryData& q, pt::ptree& tree) { return Status(0, "OK"); } +Status serializeQueryData(const QueryData& q, + const ColumnNames& cols, + pt::ptree& tree) { + for (const auto& r : q) { + pt::ptree serialized; + auto s = serializeRow(r, cols, serialized); + if (!s.ok()) { + return s; + } + tree.push_back(std::make_pair("", serialized)); + } + return Status(0, "OK"); +} + Status serializeQueryDataJSON(const QueryData& q, std::string& json) { pt::ptree tree; auto status = serializeQueryData(q, tree); diff --git a/osquery/database/tests/results_tests.cpp b/osquery/database/tests/results_tests.cpp index 71cf584e..44820785 100644 --- a/osquery/database/tests/results_tests.cpp +++ b/osquery/database/tests/results_tests.cpp @@ -69,6 +69,16 @@ TEST_F(ResultsTests, test_serialize_query_data) { EXPECT_EQ(results.first, tree); } +TEST_F(ResultsTests, test_serialize_query_data_in_column_order) { + auto results = getSerializedQueryDataWithColumnOrder(); + auto column_names = getSerializedRowColumnNames(true); + pt::ptree tree; + auto s = serializeQueryData(results.second, column_names, tree); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(s.toString(), "OK"); + EXPECT_EQ(results.first, tree); +} + TEST_F(ResultsTests, test_serialize_query_data_json) { auto results = getSerializedQueryDataJSON(); std::string json; diff --git a/osquery/distributed/distributed.cpp b/osquery/distributed/distributed.cpp index ff7e874b..0066fa41 100644 --- a/osquery/distributed/distributed.cpp +++ b/osquery/distributed/distributed.cpp @@ -92,7 +92,7 @@ Status Distributed::serializeResults(std::string& json) { pt::ptree statuses; for (const auto& result : results_) { pt::ptree qd; - auto s = serializeQueryData(result.results, qd); + auto s = serializeQueryData(result.results, result.columns, qd); if (!s.ok()) { return s; } @@ -131,7 +131,8 @@ Status Distributed::runQueries() { << sql.getMessageString(); } - DistributedQueryResult result(request, sql.rows(), sql.getStatus()); + DistributedQueryResult result( + request, sql.rows(), sql.columns(), sql.getStatus()); addResult(result); } return flushCompleted(); @@ -270,7 +271,7 @@ Status serializeDistributedQueryResult(const DistributedQueryResult& r, } pt::ptree results; - s = serializeQueryData(r.results, results); + s = serializeQueryData(r.results, r.columns, results); if (!s.ok()) { return s; } diff --git a/osquery/sql/sql.cpp b/osquery/sql/sql.cpp index 7c537849..fd1a7f3e 100644 --- a/osquery/sql/sql.cpp +++ b/osquery/sql/sql.cpp @@ -23,13 +23,25 @@ FLAG(int32, value_max, 512, "Maximum returned row value size"); CREATE_LAZY_REGISTRY(SQLPlugin, "sql"); SQL::SQL(const std::string& q) { - status_ = query(q, results_); + TableColumns table_columns; + q_ = q; + status_ = getQueryColumns(q_, table_columns); + if (status_.ok()) { + for (auto c : table_columns) { + columns_.push_back(std::get<0>(c)); + } + status_ = query(q_, results_); + } } const QueryData& SQL::rows() const { return results_; } +const ColumnNames& SQL::columns() { + return columns_; +} + bool SQL::ok() { return status_.ok(); } diff --git a/osquery/tests/test_util.cpp b/osquery/tests/test_util.cpp index b165427c..1b3e7f58 100644 --- a/osquery/tests/test_util.cpp +++ b/osquery/tests/test_util.cpp @@ -239,18 +239,41 @@ ScheduledQuery getOsqueryScheduledQuery() { return sq; } -std::pair getSerializedRow() { +ColumnNames getSerializedRowColumnNames(bool unordered_and_repeated) { + ColumnNames cn; + if (unordered_and_repeated) { + cn.push_back("repeated_column"); + } + cn.push_back("alphabetical"); + cn.push_back("foo"); + cn.push_back("meaning_of_life"); + cn.push_back("repeated_column"); + return cn; +} + +std::pair getSerializedRow(bool unordered_and_repeated) { Row r; - r["foo"] = "bar"; - r["meaning_of_life"] = "42"; + ColumnNames cns = getSerializedRowColumnNames(unordered_and_repeated); pt::ptree arr; - arr.put("foo", "bar"); - arr.put("meaning_of_life", "42"); + for (auto cn : cns) { + std::string c_value = cn + "_value"; + r[cn] = c_value; + arr.add(cn, c_value); + } return std::make_pair(arr, r); } std::pair getSerializedQueryData() { - auto r = getSerializedRow(); + auto r = getSerializedRow(false); + QueryData q = {r.second, r.second}; + pt::ptree arr; + arr.push_back(std::make_pair("", r.first)); + arr.push_back(std::make_pair("", r.first)); + return std::make_pair(arr, q); +} + +std::pair getSerializedQueryDataWithColumnOrder() { + auto r = getSerializedRow(true); QueryData q = {r.second, r.second}; pt::ptree arr; arr.push_back(std::make_pair("", r.first)); diff --git a/osquery/tests/test_util.h b/osquery/tests/test_util.h index 6b3bdeb7..71930c6c 100644 --- a/osquery/tests/test_util.h +++ b/osquery/tests/test_util.h @@ -67,15 +67,23 @@ QueryData getTestDBExpectedResults(); // need to be performed on the dataset to make the results be pair.second std::vector > getTestDBResultStream(); +// getSerializedRowColumnNames returns a vector of test column names that +// are in alphabetical order. If unordered_and_repeated is true, the +// vector includes a repeated column name and is in non-alphabetical order +ColumnNames getSerializedRowColumnNames(bool unordered_and_repeated); + // getSerializedRow() return an std::pair where pair->first is a string which // should serialize to pair->second. pair->second should deserialize // to pair->first -std::pair getSerializedRow(); +std::pair getSerializedRow(bool unordered_and_repeated = false); // getSerializedQueryData() return an std::pair where pair->first is a string // which should serialize to pair->second. pair->second should -// deserialize to pair->first +// deserialize to pair->first. getSerializedQueryDataWithColumnOrder +// returns a pair where pair->second is a tree that has a repeated column and +// the child nodes are not in alphabetical order std::pair getSerializedQueryData(); +std::pair getSerializedQueryDataWithColumnOrder(); std::pair getSerializedQueryDataJSON(); // getSerializedDiffResults() return an std::pair where pair->first is a string