Add support for additional collations in column definitions (osquery#…

…8214)
Smjert · Dec 18, 2023 · 1243c73 · 1243c73
1 parent a184da4
commit 1243c73
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 6 deletions.
diff --git a/osquery/core/sql/column.h b/osquery/core/sql/column.h
@@ -73,8 +73,10 @@ enum class ColumnOptions {
   /// This column should be hidden from '*'' selects.
   HIDDEN = 16,
 
-  // This sets the collating sequence to NOCASE
-  COLLATENOCASE = 32,
+  // Set the collation sequence for this column.
+  COLLATEBINARY = 32,
+  COLLATENOCASE = 64,
+  COLLATERTRIM = 128,
 };
 
 /// Treat column options as a set of flags.

diff --git a/osquery/core/tables.cpp b/osquery/core/tables.cpp
@@ -269,8 +269,12 @@ std::string columnDefinition(const TableColumns& columns, bool is_extension) {
     if (options & ColumnOptions::HIDDEN) {
       statement += " HIDDEN";
     }
-    if (options & ColumnOptions::COLLATENOCASE) {
+    if (options & ColumnOptions::COLLATEBINARY) {
+      statement += " COLLATE BINARY";
+    } else if (options & ColumnOptions::COLLATENOCASE) {
       statement += " COLLATE NOCASE";
+    } else if (options & ColumnOptions::COLLATERTRIM) {
+      statement += " COLLATE RTRIM";
     }
     if (i < columns.size() - 1) {
       statement += ", ";

diff --git a/specs/windows/registry.table b/specs/windows/registry.table
@@ -1,7 +1,7 @@
 table_name("registry")
 description("All of the Windows registry hives.")
 schema([
-    Column("key", TEXT, "Name of the key to search for", additional=True, collate_nocase=True),
+    Column("key", TEXT, "Name of the key to search for", additional=True, collate="nocase"),
     Column("path", TEXT, "Full path to the value", index=True),
     Column("name", TEXT, "Name of the registry value entry"),
     Column("type", TEXT, "Type of the registry value, or 'subkey' if item is a subkey"),

diff --git a/tools/codegen/gentable.py b/tools/codegen/gentable.py
@@ -73,9 +73,17 @@ def __repr__(self):
     "required": "REQUIRED",
     "optimized": "OPTIMIZED",
     "hidden": "HIDDEN",
-    "collate_nocase": "COLLATENOCASE",
+    "collate": "COLLATE",
 }
 
+# Available collation sequences to set on column definitions.
+# This should mimic the collation sequences in ColumnOptions.
+COLLATIONS = [
+    "BINARY",
+    "NOCASE",
+    "RTRIM",
+]
+
 # Column options that render tables uncacheable.
 NON_CACHEABLE = [
     "REQUIRED",
@@ -215,7 +223,17 @@ def generate(self, path, template="default"):
             for option in column.options:
                 # Only allow explicitly-defined options.
                 if option in COLUMN_OPTIONS:
-                    column_options.append("ColumnOptions::" + COLUMN_OPTIONS[option])
+                    if option == "collate":
+                        collation = str(column.options["collate"]).upper()
+                        if collation in COLLATIONS:
+                            column_options.append("ColumnOptions::" + COLUMN_OPTIONS[option] + collation)
+                        else:
+                            print(lightred(
+                                "Table %s column %s contains an unknown collation sequence: %s" % (
+                                self.table_name, column.name, column.options["collate"])))
+                            exit(1)
+                    else:
+                        column_options.append("ColumnOptions::" + COLUMN_OPTIONS[option])
                     all_options.append(COLUMN_OPTIONS[option])
                 else:
                     print(yellow(