Include messages in distributed results (#115)

osquery · Oct 6, 2023 · d6f325f · d6f325f
1 parent 2e4891a
commit d6f325f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git a/plugin/distributed/distributed.go b/plugin/distributed/distributed.go
@@ -48,6 +48,8 @@ type Result struct {
 	Rows []map[string]string `json:"rows"`
 	// QueryStats are the stats about the execution of the given query
 	QueryStats *Stats `json:"stats"`
+	// Message is the message string indicating the status of the query
+	Message string `json:"message"`
 }
 
 // WriteResultsFunc writes the results of the executed distributed queries. The
@@ -143,6 +145,7 @@ type ResultsStruct struct {
 	Queries  map[string][]map[string]string `json:"queries"`
 	Statuses map[string]OsqueryInt          `json:"statuses"`
 	Stats    map[string]Stats               `json:"stats"`
+	Messages map[string]string              `json:"messages"`
 }
 
 // Stats holds performance stats about the execution of a given query.
@@ -158,12 +161,14 @@ func (rs *ResultsStruct) UnmarshalJSON(buff []byte) error {
 	emptyRow := []map[string]string{}
 	rs.Queries = make(map[string][]map[string]string)
 	rs.Statuses = make(map[string]OsqueryInt)
+	rs.Messages = make(map[string]string)
 	// Queries can be []map[string]string OR an empty string
 	// so we need to deal with an interface to accommodate two types
 	intermediate := struct {
 		Queries  map[string]interface{} `json:"queries"`
 		Statuses map[string]OsqueryInt  `json:"statuses"`
 		Stats    map[string]Stats       `json:"stats"`
+		Messages map[string]string      `json:"messages"`
 	}{}
 	if err := json.Unmarshal(buff, &intermediate); err != nil {
 		return err
@@ -192,8 +197,9 @@ func (rs *ResultsStruct) UnmarshalJSON(buff []byte) error {
 			return fmt.Errorf("results for %q unknown type", queryName)
 		}
 	}
-	// Stats don't require any format changes
+	// Stats and messages don't require any format changes
 	rs.Stats = intermediate.Stats
+	rs.Messages = intermediate.Messages
 	return nil
 }
 
@@ -209,6 +215,9 @@ func (rs *ResultsStruct) toResults() ([]Result, error) {
 		if stats, ok := rs.Stats[queryName]; ok {
 			result.QueryStats = &stats
 		}
+		if msg, ok := rs.Messages[queryName]; ok {
+			result.Message = msg
+		}
 		results = append(results, result)
 	}
 	return results, nil

diff --git a/plugin/distributed/distributed_test.go b/plugin/distributed/distributed_test.go
@@ -63,9 +63,9 @@ func TestDistributedPlugin(t *testing.T) {
 	// Ensure correct ordering for comparison
 	sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName })
 	assert.Equal(t, []Result{
-		{"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}},
-		{"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}},
-		{"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}},
+		{"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}, ""},
+		{"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}, ""},
+		{"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}, ""},
 	},
 		results)
 
@@ -78,9 +78,24 @@ func TestDistributedPlugin(t *testing.T) {
 	// Ensure correct ordering for comparison
 	sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName })
 	assert.Equal(t, []Result{
-		{"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, nil},
-		{"query2", 0, []map[string]string{{"version": "2.4.0"}}, nil},
-		{"query3", 1, []map[string]string{}, nil},
+		{"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, nil, ""},
+		{"query2", 0, []map[string]string{{"version": "2.4.0"}}, nil, ""},
+		{"query3", 1, []map[string]string{}, nil, ""},
+	},
+		results)
+
+	// Call writeResults -- with message included
+	getCalled = false
+	resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": `{"queries":{"query1":[{"iso_8601":"2017-07-10T22:08:40Z"}],"query2":[{"version":"2.4.0"}]},"statuses":{"query1":"0","query2":"0","query3":"1"}, "stats":{"query1":{"wall_time_ms": 1, "user_time": 1, "system_time": 1, "memory": 1},"query2":{"wall_time_ms": 2, "user_time": 2, "system_time": 2, "memory": 2},"query3":{"wall_time_ms": 3, "user_time": 3, "system_time": 3, "memory": 3}}, "messages": {"query3": "distributed query is denylisted"}}`})
+	assert.False(t, getCalled)
+	assert.True(t, writeCalled)
+	assert.Equal(t, &StatusOK, resp.Status)
+	// Ensure correct ordering for comparison
+	sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName })
+	assert.Equal(t, []Result{
+		{"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}, ""},
+		{"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}, ""},
+		{"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}, "distributed query is denylisted"},
 	},
 		results)
 }