Expose table column usage in query context

plugin/table/table.go

@@ -182,6 +182,30 @@ type QueryContext struct {
 	// Constraints is a map from column name to the details of the
 	// constraints on that column.
 	Constraints map[string]ConstraintList
+
+	// ColumnsUsed contains the columns osquery reports as needed by the query.
+	// If this is nil, the extension should assume all columns may be needed.
+	ColumnsUsed []string
+
+	// ColumnsUsedBitset is the raw SQLite column-use bitset passed through by
+	// osquery. Consumers that need name-based checks should prefer ColumnsUsed.
+	ColumnsUsedBitset *uint64
+}
+
+func (q QueryContext) HasColumnUsage() bool {
+	return q.ColumnsUsed != nil || q.ColumnsUsedBitset != nil
+}
+
+func (q QueryContext) IsColumnUsed(name string) bool {
+	if q.ColumnsUsed == nil {
+		return true
+	}
+	for _, column := range q.ColumnsUsed {
+		if column == name {
+			return true
+		}
+	}
+	return false
 }
 
 // ConstraintList contains the details of the constraints for the given column.
@@ -226,7 +250,9 @@ const (
 // The following types and functions exist for parsing of the queryContext
 // JSON and are not made public.
 type queryContextJSON struct {
-	Constraints []constraintListJSON `json:"constraints"`
+	Constraints       []constraintListJSON `json:"constraints"`
+	ColumnsUsed       []string             `json:"colsUsed"`
+	ColumnsUsedBitset *uint64              `json:"colsUsedBitset"`
 }
 
 type constraintListJSON struct {
@@ -243,7 +269,11 @@ func parseQueryContext(ctxJSON string) (*QueryContext, error) {
 		return nil, errors.Wrap(err, "unmarshaling context JSON")
 	}
 
-	ctx := QueryContext{map[string]ConstraintList{}}
+	ctx := QueryContext{
+		Constraints:       map[string]ConstraintList{},
+		ColumnsUsed:       parsed.ColumnsUsed,
+		ColumnsUsedBitset: parsed.ColumnsUsedBitset,
+	}
 	for _, cList := range parsed.Constraints {
 		constraints, err := parseConstraintList(cList.List)
 		if err != nil {

plugin/table/table_test.go

@@ -57,7 +57,7 @@ func TestTablePlugin(t *testing.T) {
 
 	// Call with good action and context
 	resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "generate", "context": "{}"})
-	assert.Equal(t, QueryContext{map[string]ConstraintList{}}, calledQueryCtx)
+	assert.Equal(t, QueryContext{Constraints: map[string]ConstraintList{}}, calledQueryCtx)
 	assert.Equal(t, &StatusOK, resp.Status)
 	assert.Equal(t, osquery.ExtensionPluginResponse{
 		{
@@ -189,7 +189,7 @@ func TestParseQueryContext(t *testing.T) {
     }
   ]
 }`,
-			context: QueryContext{map[string]ConstraintList{
+			context: QueryContext{Constraints: map[string]ConstraintList{
 				"big_int": ConstraintList{ColumnTypeBigInt, []Constraint{}},
 				"double":  ConstraintList{ColumnTypeDouble, []Constraint{}},
 				"integer": ConstraintList{ColumnTypeInteger, []Constraint{}},
@@ -233,7 +233,7 @@ func TestParseQueryContext(t *testing.T) {
   ]
 }
 `,
-			context: QueryContext{map[string]ConstraintList{
+			context: QueryContext{Constraints: map[string]ConstraintList{
 				"big_int": ConstraintList{ColumnTypeBigInt, []Constraint{}},
 				"double":  ConstraintList{ColumnTypeDouble, []Constraint{{OperatorGreaterThanOrEquals, "3.1"}}},
 				"integer": ConstraintList{ColumnTypeInteger, []Constraint{}},
@@ -318,3 +318,41 @@ func TestParseVaryingQueryContexts(t *testing.T) {
 		})
 	}
 }
+
+func TestParseQueryContextColumnsUsed(t *testing.T) {
+	context, err := parseQueryContext(`{
+  "constraints": [],
+  "colsUsed": ["source", "session_id"],
+  "colsUsedBitset": 3
+}`)
+	require.NoError(t, err)
+
+	require.NotNil(t, context)
+	require.True(t, context.HasColumnUsage())
+	assert.Equal(t, []string{"source", "session_id"}, context.ColumnsUsed)
+	require.NotNil(t, context.ColumnsUsedBitset)
+	assert.Equal(t, uint64(3), *context.ColumnsUsedBitset)
+	assert.True(t, context.IsColumnUsed("source"))
+	assert.False(t, context.IsColumnUsed("text"))
+}
+
+func TestParseQueryContextEmptyColumnsUsed(t *testing.T) {
+	context, err := parseQueryContext(`{
+  "constraints": [],
+  "colsUsed": [],
+  "colsUsedBitset": 0
+}`)
+	require.NoError(t, err)
+
+	require.NotNil(t, context)
+	require.True(t, context.HasColumnUsage())
+	assert.NotNil(t, context.ColumnsUsed)
+	assert.False(t, context.IsColumnUsed("text"))
+}
+
+func TestQueryContextIsColumnUsedFallback(t *testing.T) {
+	context := QueryContext{Constraints: map[string]ConstraintList{}}
+
+	assert.False(t, context.HasColumnUsage())
+	assert.True(t, context.IsColumnUsed("text"))
+}