Merge pull request #1675 from theopolis/planner_or

Fix constraints stacking
This commit is contained in:
Teddy Reed 2015-11-24 12:25:15 -08:00
commit 204b16a946
5 changed files with 258 additions and 57 deletions

View File

@ -11,6 +11,7 @@
#include <gtest/gtest.h>
#include <osquery/core.h>
#include <osquery/logger.h>
#include <osquery/registry.h>
#include <osquery/sql.h>
@ -73,9 +74,125 @@ TEST_F(VirtualTableTests, test_sqlite3_table_joins) {
QueryData results;
// Run a query with a join within.
std::string statement =
"SELECT p.pid FROM osquery_info oi, processes p WHERE oi.pid=p.pid";
"SELECT p.pid FROM osquery_info oi, processes p WHERE oi.pid = p.pid";
auto status = queryInternal(statement, results, dbc.db());
EXPECT_TRUE(status.ok());
EXPECT_EQ(results.size(), 1U);
}
class pTablePlugin : public TablePlugin {
private:
TableColumns columns() const {
return {
{"x", INTEGER_TYPE}, {"y", INTEGER_TYPE},
};
}
public:
QueryData generate(QueryContext&) {
return {
{{"x", "1"}, {"y", "2"}}, {{"x", "2"}, {"y", "1"}},
};
}
private:
FRIEND_TEST(VirtualTableTests, test_constraints_stacking);
};
class kTablePlugin : public TablePlugin {
private:
TableColumns columns() const {
return {
{"x", INTEGER_TYPE}, {"z", INTEGER_TYPE},
};
}
public:
QueryData generate(QueryContext&) {
return {
{{"x", "1"}, {"z", "2"}}, {{"x", "2"}, {"z", "1"}},
};
}
private:
FRIEND_TEST(VirtualTableTests, test_constraints_stacking);
};
TEST_F(VirtualTableTests, test_constraints_stacking) {
// Add two testing tables to the registry.
Registry::add<pTablePlugin>("table", "p");
Registry::add<kTablePlugin>("table", "k");
auto dbc = SQLiteDBManager::get();
{
// To simplify the attach, just access the column definition directly.
auto p = std::make_shared<pTablePlugin>();
attachTableInternal("p", p->columnDefinition(), dbc.db());
auto k = std::make_shared<kTablePlugin>();
attachTableInternal("k", k->columnDefinition(), dbc.db());
}
QueryData results;
std::string statement;
std::map<std::string, std::string> expected;
statement = "select k.x from (select * from k) k2, p, k where k.x = p.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement =
"select k.x from (select * from k where z = 1) k2, p, k where k.x = p.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement = "select k.x from k k1, (select * from p) p1, k where k.x = p1.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
statement =
"select k.x from (select * from p) p1, k, (select * from k) k2 where k.x "
"= p1.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement =
"select k.x from (select * from p) p1, k, (select * from k where z = 2) "
"k2 where k.x = p1.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement =
"select k.x from k, (select * from p) p1, k k2, (select * from k where z "
"= 1) k3 where k.x = p1.x";
expected = {{"k.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement =
"select p.x from (select * from k where z = 1) k1, (select * from k "
"where z != 1) k2, p where p.x = k2.x";
expected = {{"p.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
statement =
"select p.x from (select * from k, (select x as xx from k where x = 1) "
"k2 where z = 1) k1, (select * from k where z != 1) k2, p, k as k3 where "
"p.x = k2.x";
expected = {{"p.x", "1"}};
queryInternal(statement, results, dbc.db());
EXPECT_EQ(results[0], expected);
results.clear();
}
}

View File

@ -16,10 +16,9 @@ namespace osquery {
namespace tables {
namespace sqlite {
int xOpen(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor) {
int xOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
int rc = SQLITE_NOMEM;
BaseCursor *pCur = new BaseCursor;
if (pCur) {
memset(pCur, 0, sizeof(BaseCursor));
*ppCursor = (sqlite3_vtab_cursor *)pCur;
@ -31,19 +30,27 @@ int xOpen(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor) {
int xClose(sqlite3_vtab_cursor *cur) {
BaseCursor *pCur = (BaseCursor *)cur;
const auto *pVtab = (VirtualTable *)cur->pVtab;
if (pVtab != nullptr) {
// Reset all constraints for the virtual table content.
if (pVtab->content->constraints.size() > 0) {
// As each cursor is closed remove the potential constraints it used.
// Cursors without constraints (full scans) are kept open.
pVtab->content->constraints.pop_front();
}
pVtab->content->constraints_cursor = nullptr;
pVtab->content->constraints_index = 0;
pVtab->content->current_term = -1;
}
delete pCur;
return SQLITE_OK;
}
int xEof(sqlite3_vtab_cursor *cur) {
BaseCursor *pCur = (BaseCursor *)cur;
auto *pVtab = (VirtualTable *)cur->pVtab;
if (pCur->row >= pVtab->content->n) {
if (pCur->row >= pCur->n) {
// If the requested row exceeds the size of the row set then all rows
// have been visited, clear the data container.
QueryData().swap(pVtab->content->data);
return true;
}
return false;
@ -75,7 +82,6 @@ int xCreate(sqlite3 *db,
sqlite3_vtab **ppVtab,
char **pzErr) {
auto *pVtab = new VirtualTable;
if (!pVtab || argc == 0 || argv[0] == nullptr) {
delete pVtab;
return SQLITE_NOMEM;
@ -84,9 +90,9 @@ int xCreate(sqlite3 *db,
memset(pVtab, 0, sizeof(VirtualTable));
pVtab->content = new VirtualTableContent;
// Create a TablePlugin Registry call, expect column details as the response.
PluginResponse response;
pVtab->content->name = std::string(argv[0]);
// Get the table column information.
auto status = Registry::call(
"table", pVtab->content->name, {{"action", "columns"}}, response);
@ -96,6 +102,7 @@ int xCreate(sqlite3 *db,
return SQLITE_ERROR;
}
// Generate an SQL create table statement from the retrieved column details.
auto statement =
"CREATE TABLE " + pVtab->content->name + columnDefinition(response);
int rc = sqlite3_declare_vtab(db, statement.c_str());
@ -105,19 +112,19 @@ int xCreate(sqlite3 *db,
return (rc != SQLITE_OK) ? rc : SQLITE_ERROR;
}
// Keep a local copy of the column details in the VirtualTableContent struct.
// This allows introspection into the column type without additional calls.
for (const auto &column : response) {
pVtab->content->columns.push_back(
std::make_pair(column.at("name"), columnTypeName(column.at("type"))));
}
*ppVtab = (sqlite3_vtab *)pVtab;
return rc;
}
int xColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int col) {
const BaseCursor *pCur = (BaseCursor *)cur;
BaseCursor *pCur = (BaseCursor *)cur;
const auto *pVtab = (VirtualTable *)cur->pVtab;
if (col >= static_cast<int>(pVtab->content->columns.size())) {
// Requested column index greater than column set size.
return SQLITE_ERROR;
@ -125,13 +132,13 @@ int xColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int col) {
const auto &column_name = pVtab->content->columns[col].first;
const auto &type = pVtab->content->columns[col].second;
if (pCur->row >= pVtab->content->data.size()) {
if (pCur->row >= pCur->data.size()) {
// Request row index greater than row set size.
return SQLITE_ERROR;
}
// Attempt to cast each xFilter-populated row/column to the SQLite type.
const auto &value = pVtab->content->data[pCur->row][column_name];
const auto &value = pCur->data[pCur->row][column_name];
if (type == TEXT_TYPE) {
sqlite3_result_text(ctx, value.c_str(), value.size(), SQLITE_STATIC);
} else if (type == INTEGER_TYPE) {
@ -167,28 +174,60 @@ int xColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int col) {
static int xBestIndex(sqlite3_vtab *tab, sqlite3_index_info *pIdxInfo) {
auto *pVtab = (VirtualTable *)tab;
ConstraintSet().swap(pVtab->content->constraints);
int expr_index = 0;
int cost = 0;
ConstraintSet constraints;
// Keep track of the index used for each valid constraint.
// Expect this index to correspond with argv within xFilter.
size_t expr_index = 0;
// If any constraints are unusable increment the cost of the index.
size_t cost = 0;
// Expressions operating on the same virtual table are loosely identified by
// the consecutive sets of terms each of the constraint sets are applied onto.
// Subsequent attempts from failed (unusable) constraints replace the set,
// while new sets of terms append.
int term = -1;
if (pIdxInfo->nConstraint > 0) {
for (size_t i = 0; i < static_cast<size_t>(pIdxInfo->nConstraint); ++i) {
// Record the term index (this index exists across all expressions).
term = pIdxInfo->aConstraint[i].iTermOffset;
if (!pIdxInfo->aConstraint[i].usable) {
// A higher cost less priority, prefer more usable query constraints.
cost += 10;
// TODO: OR is not usable.
continue;
}
// Lookup the column name given an index into the table column set.
const auto &name =
pVtab->content->columns[pIdxInfo->aConstraint[i].iColumn].first;
pVtab->content->constraints.push_back(
// Save a pair of the name and the constraint operator.
// Use this constraint during xFilter by performing a scan and column
// name lookup through out all cursor constraint lists.
constraints.push_back(
std::make_pair(name, Constraint(pIdxInfo->aConstraint[i].op)));
pIdxInfo->aConstraintUsage[i].argvIndex = ++expr_index;
}
}
// Set the estimated cost based on the number of unusable terms.
pIdxInfo->estimatedCost = cost;
if (cost == 0 && term != -1) {
// This set of constraints is 100% usable.
// Add the constraint set to the table's tracked constraints.
if (term != pVtab->content->current_term ||
pVtab->content->constraints.size() == 0) {
pVtab->content->constraints.push_back(constraints);
} else {
// It is possible this strategy is successful after a set of failures.
// If there was a previous failure, replace the constraints.
pVtab->content->constraints.back() = constraints;
}
} else {
// Failed.
if (term != -1 && term != pVtab->content->current_term) {
// If this failure is on another set of terms, add empty constraints.
pVtab->content->constraints.push_back(ConstraintSet());
}
}
pVtab->content->current_term = term;
return SQLITE_OK;
}
@ -199,45 +238,61 @@ static int xFilter(sqlite3_vtab_cursor *pVtabCursor,
sqlite3_value **argv) {
BaseCursor *pCur = (BaseCursor *)pVtabCursor;
auto *pVtab = (VirtualTable *)pVtabCursor->pVtab;
auto *content = pVtab->content;
pCur->row = 0;
pVtab->content->n = 0;
pCur->n = 0;
QueryContext context;
for (size_t i = 0; i < pVtab->content->columns.size(); ++i) {
for (size_t i = 0; i < content->columns.size(); ++i) {
// Set the column affinity for each optional constraint list.
// There is a separate list for each column name.
context.constraints[pVtab->content->columns[i].first].affinity =
pVtab->content->columns[i].second;
context.constraints[content->columns[i].first].affinity =
content->columns[i].second;
}
// Filtering between cursors happens iteratively, not consecutively.
// If there are multiple sets of constraints, they apply to each cursor.
if (content->constraints_cursor == nullptr) {
content->constraints_cursor = pVtabCursor;
} else if (content->constraints_cursor != pVtabCursor) {
content->constraints_index += 1;
if (content->constraints_index >= content->constraints.size()) {
content->constraints_index = 0;
}
content->constraints_cursor = pVtabCursor;
}
// Iterate over every argument to xFilter, filling in constraint values.
if (argc > 0) {
for (size_t i = 0; i < static_cast<size_t>(argc); ++i) {
auto expr = (const char *)sqlite3_value_text(argv[i]);
if (expr == nullptr) {
// SQLite did not expose the expression value.
continue;
if (content->constraints.size() > 0) {
auto &constraints = content->constraints[content->constraints_index];
if (argc > 0) {
for (size_t i = 0; i < static_cast<size_t>(argc); ++i) {
auto expr = (const char *)sqlite3_value_text(argv[i]);
if (expr == nullptr || expr[0] == 0) {
// SQLite did not expose the expression value.
continue;
}
// Set the expression from SQLite's now-populated argv.
auto &constraint = constraints[i];
constraint.second.expr = std::string(expr);
// Add the constraint to the column-sorted query request map.
context.constraints[constraint.first].add(constraint.second);
}
// Set the expression from SQLite's now-populated argv.
pVtab->content->constraints[i].second.expr = std::string(expr);
// Add the constraint to the column-sorted query request map.
context.constraints[pVtab->content->constraints[i].first].add(
pVtab->content->constraints[i].second);
} else if (constraints.size() > 0) {
// Constraints failed.
}
}
// Reset the virtual table contents.
pVtab->content->data.clear();
pCur->data.clear();
// Generate the row data set.
PluginRequest request = {{"action", "generate"}};
TablePlugin::setRequestFromContext(context, request);
Registry::call("table", pVtab->content->name, request, pVtab->content->data);
Registry::call("table", pVtab->content->name, request, pCur->data);
// Set the number of rows.
pVtab->content->n = pVtab->content->data.size();
pCur->n = pCur->data.size();
return SQLITE_OK;
}
}

View File

@ -10,6 +10,8 @@
#pragma once
#include <deque>
#include <osquery/tables.h>
#include "osquery/core/conversions.h"
@ -17,6 +19,34 @@
namespace osquery {
/**
* @brief osquery virtual table connection.
*
* This object is the SQLite database's virtual table context.
* When the virtual table is created/connected the name and columns are
* retrieved via the TablePlugin call API. The details are kept in this context
* so column parsing and row walking does not require additional Registry calls.
*
* When tables are accessed as the result of an SQL statement a QueryContext is
* created to represent metadata that can be used by the virtual table
* implementation code. Thus the code that generates rows can choose to emit
* additional data, restrict based on constraints, or potentially yield from
* a cache or choose not to generate certain columns.
*/
struct VirtualTableContent {
/// Friendly name for the table.
TableName name;
/// Table column structure, retrieved once via the TablePlugin call API.
TableColumns columns;
/// Transient set of virtual table access constraints.
std::deque<ConstraintSet> constraints;
/// Index into the list of constraints.
sqlite3_vtab_cursor *constraints_cursor{nullptr};
size_t constraints_index{0};
/// Last term successfully parsed by xBestIndex.
int current_term{-1};
};
/**
* @brief osquery cursor object.
*
@ -25,16 +55,12 @@ namespace osquery {
struct BaseCursor {
/// SQLite virtual table cursor.
sqlite3_vtab_cursor base;
/// Current cursor position.
size_t row;
};
struct VirtualTableContent {
TableName name;
TableColumns columns;
/// Table data generated from last access.
QueryData data;
ConstraintSet constraints;
size_t n;
/// Current cursor position.
size_t row{0};
/// Total number of rows.
size_t n{0};
};
/**
@ -45,7 +71,7 @@ struct VirtualTableContent {
*/
struct VirtualTable {
sqlite3_vtab base;
VirtualTableContent *content;
VirtualTableContent *content{nullptr};
};
/// Attach a table plugin name to an in-memory SQLite database.

View File

@ -42,10 +42,13 @@ class AdditionalFeatureTests(test_base.ProcessGenerator, unittest.TestCase):
# Get a daemon process, loaded with the default test configuration.
# We'll add a config override (overwrite) for the "packs" key.
# THis will point a single pack at the config written above.
daemon = self._run_daemon(overwrite={
daemon = self._run_daemon({
"disable_watchdog": True,
},
overwrite={
"packs": {
"test_pack": query_pack_path
}
},
})
self.assertTrue(daemon.isAlive())

View File

@ -432,7 +432,7 @@ def flaky(gen):
except Exception as e:
exc_type, exc_obj, exc_tb = sys.exc_info()
fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
exceptions.append((str(e), fname, exc_tb.tb_lineno))
exceptions.append((e, fname, exc_tb.tb_lineno))
return False
def wrapper(this):
for i in range(3):
@ -443,9 +443,9 @@ def flaky(gen):
print("Test (attempt %d) %s::%s failed: %s (%s:%d)" % (
i,
this.__class__.__name__,
gen.__name__, exc[0], exc[1], exc[2]))
gen.__name__, str(exc[0]), exc[1], exc[2]))
i += 1
return False
raise exceptions[0][0]
return wrapper
class Tester(object):