Added new client message type FlowRequest (Velocidex#2391)

This message combines all collections in the same flow into a single
message, reducing the total number of task files in message
queues. This should help to reduce IO operations and increase
throughput

actions/events.go

@@ -224,13 +224,7 @@ func (self UpdateEventTable) Run(
 			defer service_wg.Done()
 
 			// Name of the query we are running.
-			name := ""
-			for _, q := range event.Query {
-				if q.Name != "" {
-					name = q.Name
-				}
-			}
-
+			name := GetQueryName(event.Query)
 			if name != "" {
 				logger.Info("<green>Starting</> monitoring query %s", name)
 			}

actions/events_test.go

@@ -133,8 +133,8 @@ func (self *EventsTestSuite) TestEventTableUpdate() {
 	// Only one query will be selected now since no label is set
 	// on the client.
 	assert.Equal(self.T(), len(message.UpdateEventTable.Event), 1)
-	assert.Equal(self.T(), getQueryName(message.UpdateEventTable.Event[0]),
-		"EventArtifact1")
+	assert.Equal(self.T(), actions.GetQueryName(
+		message.UpdateEventTable.Event[0].Query), "EventArtifact1")
 
 	// Set the new table, this will execute the new queries and
 	// start the new table.
@@ -248,15 +248,6 @@ func (self *EventsTestSuite) TestEventTableUpdate() {
 	assert.Contains(self.T(), string(data), "EventArtifact2")
 }
 
-func getQueryName(args *actions_proto.VQLCollectorArgs) string {
-	for _, query := range args.Query {
-		if query.Name != "" {
-			return query.Name
-		}
-	}
-	return ""
-}
-
 func TestEventsTestSuite(t *testing.T) {
 	suite.Run(t, &EventsTestSuite{})
 }

actions/utils.go

@@ -0,0 +1,12 @@
+package actions
+
+import actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
+
+func GetQueryName(args []*actions_proto.VQLRequest) string {
+	for _, query := range args {
+		if query.Name != "" {
+			return query.Name
+		}
+	}
+	return ""
+}

actions/vql.go

@@ -111,6 +111,8 @@ func (self VQLClientAction) StartQuery(
 		return
 	}
 
+	name := GetQueryName(arg.Query)
+
 	// Clients do not have a copy of artifacts so they need to be
 	// sent all artifacts from the server.
 	manager, err := services.GetRepositoryManager(config_obj)
@@ -163,7 +165,7 @@ func (self VQLClientAction) StartQuery(
 			"incorrect or missed results or even crashes.")
 	}
 
-	scope.Log("INFO:Starting query execution.")
+	scope.Log("INFO:Starting query execution for %v.", name)
 
 	throttler := NewThrottler(ctx, scope, float64(rate),
 		float64(cpu_limit), float64(iops_limit))
@@ -188,18 +190,19 @@ func (self VQLClientAction) StartQuery(
 			responder.RaiseError(ctx, msg)
 		}
 
-		scope.Log("INFO:Collection is done after %v", time.Since(start))
+		scope.Log("INFO:Collection %v is done after %v", name, time.Since(start))
 	}()
 
 	ok, err := CheckPreconditions(ctx, scope, arg)
 	if err != nil {
-		scope.Log("While evaluating preconditions: %v", err)
-		responder.RaiseError(ctx, fmt.Sprintf("While evaluating preconditions: %v", err))
+		scope.Log("%v: While evaluating preconditions: %v", name, err)
+		responder.RaiseError(ctx,
+			fmt.Sprintf("While evaluating preconditions: %v", err))
 		return
 	}
 
 	if !ok {
-		scope.Log("Skipping query due to preconditions")
+		scope.Log("INFO:%v: Skipping query due to preconditions", name)
 		responder.Return(ctx)
 		return
 	}
@@ -247,7 +250,7 @@ func (self VQLClientAction) StartQuery(
 
 			case <-time.After(time.Second * time.Duration(heartbeat)):
 				responder.Log(ctx, logging.DEFAULT,
-					fmt.Sprintf("Time %v: %s: Waiting for rows.",
+					fmt.Sprintf("%v: Time %v: %s: Waiting for rows.", name,
 						(uint64(time.Now().UTC().UnixNano()/1000)-
 							query_start)/1000000, query.Name))
 
@@ -279,7 +282,8 @@ func (self VQLClientAction) StartQuery(
 					responder.Log(ctx,
 						logging.DEFAULT,
 						fmt.Sprintf(
-							"Time %v: %s: Sending response part %d %s (%d rows).",
+							"%v: Time %v: %s: Sending response part %d %s (%d rows).",
+							name,
 							(response.Timestamp-query_start)/1000000,
 							query.Name,
 							result.Part,
@@ -296,7 +300,7 @@ func (self VQLClientAction) StartQuery(
 
 	if uploader.Count > 0 {
 		responder.Log(ctx, logging.DEFAULT,
-			fmt.Sprintf("Uploaded %v files.", uploader.Count))
+			fmt.Sprintf("%v: Uploaded %v files.", name, uploader.Count))
 	}
 
 	responder.Return(ctx)

executor/executor.go

@@ -152,34 +152,30 @@ func (self *ClientExecutor) processRequestPlugin(
 		return
 	}
 
-	// Handle the requests. This used to be a plugin registration
-	// process but there are very few plugins any more and so it
-	// is easier to hard code this.
-	responder_obj := responder.NewResponder(ctx, config_obj, req, self.Outbound)
-	defer responder_obj.Close()
-
-	// Each request has its own context.
-	query_ctx, closer := flow_manager.FlowContext(req).NewQueryContext(
-		responder_obj, req)
-	defer closer()
-
+	// This is the old deprecated VQLClientAction that is sent for old
+	// client compatibility. New clients ignore this and only process
+	// a FlowRequest message.
 	if req.VQLClientAction != nil {
-		// Control concurrency on the executor only.
-		if !req.Urgent {
-			cancel, err := self.concurrency.StartConcurrencyControl(query_ctx)
-			if err != nil {
-				responder_obj.RaiseError(query_ctx, fmt.Sprintf("%v", err))
-				return
-			}
-			defer cancel()
-		}
+		return
+	}
 
-		actions.VQLClientAction{}.StartQuery(
-			config_obj, query_ctx, responder_obj, req.VQLClientAction)
+	if req.FlowRequest != nil {
+		self.ProcessFlowRequest(ctx, config_obj, req)
 		return
 	}
 
 	if req.UpdateEventTable != nil {
+		// Handle the requests. This used to be a plugin registration
+		// process but there are very few plugins any more and so it
+		// is easier to hard code this.
+		responder_obj := responder.NewResponder(ctx, config_obj, req, self.Outbound)
+		defer responder_obj.Close()
+
+		// Each request has its own context.
+		flow_context := flow_manager.FlowContext(req)
+		query_ctx, closer := flow_context.NewQueryContext(responder_obj)
+		defer closer()
+
 		actions.UpdateEventTable{}.Run(
 			config_obj, query_ctx, responder_obj, req.UpdateEventTable)
 		return

executor/executor_test.go

@@ -130,11 +130,12 @@ func (self *ExecutorTestSuite) TestLogMessages() {
 	executor.Inbound <- &crypto_proto.VeloMessage{
 		AuthState: crypto_proto.VeloMessage_AUTHENTICATED,
 		SessionId: flow_id,
-		VQLClientAction: &actions_proto.VQLCollectorArgs{
-			Query: []*actions_proto.VQLRequest{
-				// Log 100 messages
-				{VQL: "SELECT log(message='log %v', args=count()) FROM range(end=10)"},
-			},
+		FlowRequest: &crypto_proto.FlowRequest{
+			VQLClientActions: []*actions_proto.VQLCollectorArgs{{
+				Query: []*actions_proto.VQLRequest{
+					// Log 100 messages
+					{VQL: "SELECT log(message='log %v', args=count()) FROM range(end=10)"},
+				}}},
 		},
 		RequestId: 1}
 

executor/flows.go

@@ -0,0 +1,46 @@
+package executor
+
+import (
+	"context"
+	"fmt"
+
+	"www.velocidex.com/golang/velociraptor/actions"
+	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
+	crypto_proto "www.velocidex.com/golang/velociraptor/crypto/proto"
+	"www.velocidex.com/golang/velociraptor/responder"
+)
+
+func (self *ClientExecutor) ProcessFlowRequest(
+	ctx context.Context,
+	config_obj *config_proto.Config, req *crypto_proto.VeloMessage) {
+
+	flow_manager := responder.GetFlowManager(ctx, config_obj)
+	flow_context := flow_manager.FlowContext(req)
+
+	// Control concurrency for the entire collection at once.
+	if !req.Urgent {
+		cancel, err := self.concurrency.StartConcurrencyControl(ctx)
+		if err != nil {
+			responder_obj := responder.NewResponder(ctx, config_obj,
+				req, self.Outbound)
+			defer responder_obj.Close()
+
+			responder_obj.RaiseError(ctx, fmt.Sprintf("%v", err))
+			return
+		}
+		defer cancel()
+	}
+
+	for _, arg := range req.FlowRequest.VQLClientActions {
+		responder_obj := responder.NewResponder(ctx, config_obj,
+			req, self.Outbound)
+		defer responder_obj.Close()
+
+		// Each request has its own context.
+		query_ctx, closer := flow_context.NewQueryContext(responder_obj)
+		defer closer()
+
+		actions.VQLClientAction{}.StartQuery(
+			config_obj, query_ctx, responder_obj, arg)
+	}
+}

flows/artifacts.go

@@ -739,7 +739,7 @@ func (self *FlowRunner) ProcessSingleMessage(
 					Cancel:    &crypto_proto.Cancel{},
 					SessionId: job.SessionId,
 				},
-				true /* notify */, nil)
+				services.NOTIFY_CLIENT, utils.BackgroundWriter)
 			if err != nil {
 				return fmt.Errorf("Queueing for client %v: %w",
 					job.Source, err)

flows/artifacts_test.go

@@ -247,9 +247,16 @@ func (self *TestSuite) TestResourceLimits() {
 		context.Background(), self.client_id)
 	assert.NoError(self.T(), err)
 
-	// Two requests since there are two source preconditions on
-	// Generic.Client.Info
-	assert.Equal(self.T(), len(messages), 2)
+	// The Generic.Client.Info has two source conditions so it
+	// contains two queries. To maintain backwards compatibility with
+	// older clients, GetClientTasks should have two old style
+	// VQLClientAction request and a new FlowRequest. Old clients will
+	// ignore the old requests and new clients will ignore the old
+	// style requests.
+	assert.Equal(self.T(), len(messages), 3)
+	assert.True(self.T(), messages[0].VQLClientAction != nil)
+	assert.True(self.T(), messages[1].VQLClientAction != nil)
+	assert.Equal(self.T(), len(messages[2].FlowRequest.VQLClientActions), 2)
 
 	// Send one row.
 	message := &crypto_proto.VeloMessage{

flows/housekeeping.go

@@ -28,6 +28,7 @@ import (
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/services"
+	utils "www.velocidex.com/golang/velociraptor/utils"
 )
 
 var (
@@ -79,7 +80,8 @@ func CheckClientStatus(
 
 		clientEventUpdateCounter.Inc()
 		err := client_manager.QueueMessageForClient(
-			ctx, client_id, update_message, true, nil)
+			ctx, client_id, update_message,
+			services.NOTIFY_CLIENT, utils.BackgroundWriter)
 		if err != nil {
 			return err
 		}

flows/limits.go

@@ -11,6 +11,7 @@ import (
 	crypto_proto "www.velocidex.com/golang/velociraptor/crypto/proto"
 	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
 	"www.velocidex.com/golang/velociraptor/services"
+	utils "www.velocidex.com/golang/velociraptor/utils"
 )
 
 var (
@@ -78,5 +79,6 @@ func cancelCollection(
 		&crypto_proto.VeloMessage{
 			Cancel:    &crypto_proto.Cancel{},
 			SessionId: flow_id,
-		}, true /* notify */, nil)
+		},
+		services.NOTIFY_CLIENT, utils.BackgroundWriter)
 }

responder/flow_manager.go

@@ -166,8 +166,7 @@ func (self *FlowContext) NextUploadId() int64 {
 	return new_id - 1
 }
 
-func (self *FlowContext) NewQueryContext(
-	responder *Responder, req *crypto_proto.VeloMessage) (
+func (self *FlowContext) NewQueryContext(responder *Responder) (
 	ctx context.Context, closer func()) {
 	self.mu.Lock()
 	defer self.mu.Unlock()
@@ -178,7 +177,6 @@ func (self *FlowContext) NewQueryContext(
 	result := &QueryContext{
 		flow_id: self.flow_id,
 		cancel:  cancel,
-		req:     req,
 		id:      utils.GetId(),
 	}
 

server/server_test.go

@@ -407,7 +407,10 @@ func (self *ServerTestSuite) TestScheduleCollection() {
 
 	tasks, err := client_info_manager.PeekClientTasks(context.Background(), self.client_id)
 	assert.NoError(t, err)
-	assert.Equal(t, len(tasks), 2)
+	assert.Equal(t, len(tasks), 1)
+
+	// The request sends a single FlowRequest task with two queries
+	assert.Equal(t, len(tasks[0].FlowRequest.VQLClientActions), 2)
 
 	collection_context := &flows_proto.ArtifactCollectorContext{}
 	path_manager := paths.NewFlowPathManager(self.client_id, flow_id)
@@ -638,7 +641,8 @@ func (self *ServerTestSuite) TestCancellation() {
 	assert.NoError(t, err)
 
 	// Generic.Client.Info has two source preconditions in parallel
-	assert.Equal(t, len(tasks), 2)
+	assert.Equal(t, len(tasks), 1)
+	assert.Equal(t, len(tasks[0].FlowRequest.VQLClientActions), 2)
 
 	// Cancelling the flow will notify the client immediately.
 	launcher, err := services.GetLauncher(self.ConfigObj)

services/client_info.go

@@ -11,6 +11,8 @@ import (
 )
 
 const (
+	NOTIFY_CLIENT = true
+
 	Unknown ClientOS = iota
 	Windows
 	Linux

services/client_info/tasks.go

@@ -287,6 +287,21 @@ func (self *ClientInfoManager) GetClientTasks(
 		if err != nil {
 			return nil, err
 		}
+
+		// Handle backwards compatibility with older clients by expanding
+		// FlowRequest into separate VQLClientActions. Newer clients will
+		// ignore bare VQLClientActions and older clients will ignore
+		// FlowRequest.
+		if message.FlowRequest != nil {
+			for _, request := range message.FlowRequest.VQLClientActions {
+				result = append(result, &crypto_proto.VeloMessage{
+					SessionId:       message.SessionId,
+					RequestId:       message.RequestId,
+					VQLClientAction: request,
+				})
+			}
+		}
+
 		result = append(result, message)
 	}