Tables with additional and no index should not use rowid

osquery · Dec 10, 2019 · fa33df3 · fa33df3
1 parent 5cf1514
commit fa33df3
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 9 deletions.
diff --git a/osquery/core/tables.cpp b/osquery/core/tables.cpp
@@ -24,6 +24,8 @@ CREATE_LAZY_REGISTRY(TablePlugin, "table");
 size_t TablePlugin::kCacheInterval = 0;
 size_t TablePlugin::kCacheStep = 0;
 
+#define kDisableRowId "WITHOUT ROWID"
+
 Status TablePlugin::addExternal(const std::string& name,
                                 const PluginResponse& response) {
   // Attach the table.
@@ -263,12 +265,12 @@ std::string columnDefinition(const TableColumns& columns, bool is_extension) {
     statement +=
         '`' + std::get<0>(column) + "` " + columnTypeName(std::get<1>(column));
     auto& options = std::get<2>(column);
-    if (options & (ColumnOptions::INDEX | ColumnOptions::ADDITIONAL)) {
-      if (options & ColumnOptions::INDEX) {
-        indexed = true;
-      }
+    if (options & ColumnOptions::INDEX) {
       pkeys.push_back(std::get<0>(column));
-      epilog["WITHOUT ROWID"] = true;
+      indexed = true;
+    }
+    if (options & (ColumnOptions::INDEX | ColumnOptions::ADDITIONAL)) {
+      epilog[kDisableRowId] = true;
     }
     if (options & ColumnOptions::HIDDEN) {
       statement += " HIDDEN";
@@ -278,10 +280,13 @@ std::string columnDefinition(const TableColumns& columns, bool is_extension) {
     }
   }
 
-  // If there are only 'additional' columns (rare), do not attempt a pkey.
-  if (!indexed) {
-    epilog["WITHOUT ROWID"] = false;
+  // If there are only 'additional' columns (rare), pkey is the 'unique row'.
+  // Otherwise an additional constraint will create duplicate rowids.
+  if (!indexed && epilog[kDisableRowId]) {
     pkeys.clear();
+    for (const auto& column : columns) {
+      pkeys.push_back(std::get<0>(column));
+    }
   }
 
   // Append the primary keys, if any were defined.
@@ -300,7 +305,7 @@ std::string columnDefinition(const TableColumns& columns, bool is_extension) {
   // keep the rowid column, as we need it to reference rows when handling UPDATE
   // and DELETE queries
   if (is_extension) {
-    epilog["WITHOUT ROWID"] = false;
+    epilog[kDisableRowId] = false;
   }
 
   statement += ')';

diff --git a/osquery/sql/tests/virtual_table.cpp b/osquery/sql/tests/virtual_table.cpp
@@ -117,6 +117,33 @@ TEST_F(VirtualTableTests, test_tableplugin_moreoptions) {
   EXPECT_EQ(expected_statement, columnDefinition(response, true, false));
 }
 
+class additionalOnlyTablePlugin : public TablePlugin {
+ private:
+  TableColumns columns() const override {
+    return {
+        std::make_tuple("id", INTEGER_TYPE, ColumnOptions::DEFAULT),
+        std::make_tuple("username", TEXT_TYPE, ColumnOptions::ADDITIONAL),
+        std::make_tuple("name", TEXT_TYPE, ColumnOptions::DEFAULT),
+    };
+  }
+
+ private:
+  FRIEND_TEST(VirtualTableTests, test_tableplugin_additionalonly);
+};
+
+TEST_F(VirtualTableTests, test_tableplugin_additionalonly) {
+  auto table = std::make_shared<additionalOnlyTablePlugin>();
+
+  PluginResponse response;
+  PluginRequest request = {{"action", "columns"}};
+  EXPECT_TRUE(table->call(request, response).ok());
+
+  std::string expected_statement =
+      "(`id` INTEGER, `username` TEXT, `name` TEXT, PRIMARY KEY (`id`, "
+      "`username`, `name`)) WITHOUT ROWID";
+  EXPECT_EQ(expected_statement, columnDefinition(response, true, false));
+}
+
 class aliasesTablePlugin : public TablePlugin {
  private:
   TableColumns columns() const override {