Bugfix: In flight flow calculation was stuck sometimes (#4668)

If flows crashed before sending any status then the server would not
update them. This resulted in repeatedly sending flow status request
without making progress.

Author: scudette

.gitignore

@@ -26,4 +26,4 @@ config/ab0x.go
 
 __debug*
 debug.test*
-enriched.json*
+artifacts/testdata/server/hunts/H.*

actions/proto/vql.proto

@@ -228,7 +228,7 @@ message ClientInfo {
     map<string, int64> in_flight_flows = 28;
 
 
-    // A list of indexed metadata fields. There are not all metadata
+    // A list of indexed metadata fields. These are not all metadata
     // fields, only the ones that are important enough to be indexed.
     map<string, string> metadata = 29;
 }

flows/client_flow_runner.go

@@ -513,31 +513,38 @@ func (self *ClientFlowRunner) handleUnknwonFlow(
 	ctx context.Context, client_id, flow_id string,
 	msg *crypto_proto.FlowStats) error {
 
-	flow_path_manager := paths.NewFlowPathManager(client_id, flow_id)
-	db, err := datastore.GetDB(self.config_obj)
+	if len(msg.QueryStatus) == 0 {
+		return nil
+	}
+
+	launcher_service, err := services.GetLauncher(self.config_obj)
 	if err != nil {
 		return err
 	}
 
-	// Just a blind write will eventually hit the disk.
-	stats := &flows_proto.ArtifactCollectorContext{}
-	err = db.GetSubject(self.config_obj, flow_path_manager.Stats(), stats)
+	// If we dont know anything about the flow, ignore it.
+	collection_context, err := launcher_service.Storage().LoadCollectionContext(
+		ctx, self.config_obj, client_id, flow_id)
 	if err != nil {
 		return nil
 	}
 
 	// Mark all the stats as terminated if they are still running.
-	for _, s := range stats.QueryStats {
-		if s.Status == crypto_proto.VeloStatus_PROGRESS {
-			s.Status = crypto_proto.VeloStatus_GENERIC_ERROR
-			s.ErrorMessage = msg.QueryStatus[0].ErrorMessage
+	if len(collection_context.QueryStats) == 0 {
+		collection_context.QueryStats = append(
+			collection_context.QueryStats, msg.QueryStatus...)
+	} else {
+		for _, s := range collection_context.QueryStats {
+			if s.Status == crypto_proto.VeloStatus_PROGRESS {
+				s.Status = msg.QueryStatus[0].Status
+				s.ErrorMessage = msg.QueryStatus[0].ErrorMessage
+			}
 		}
 	}
 
-	launcher.UpdateFlowStats(stats)
-
-	return db.SetSubjectWithCompletion(self.config_obj,
-		flow_path_manager.Stats(), stats, nil)
+	// Update the flow.
+	return launcher_service.Storage().WriteFlow(ctx, self.config_obj,
+		collection_context, utils.BackgroundWriter)
 }
 
 func (self *ClientFlowRunner) FlowStats(

flows/client_flow_runner_test.go

@@ -3,7 +3,6 @@ package flows_test
 import (
 	"context"
 	"fmt"
-	"os"
 	"regexp"
 	"strings"
 	"sync"
@@ -1025,42 +1024,62 @@ func (self *ServerTestSuite) TestCancellation() {
 
 // Test an unknown flow. What happens when the server receives a
 // message to an unknown flow.
+
+// This message is received in response to the in flight keep
+// alive. If the client crashes, the server will send a status request
+// for the in flight flow. But the client does not know about this
+// flow id. The server will terminate the flow.
 func (self *ServerTestSuite) TestUnknownFlow() {
 	t := self.T()
 
-	db, err := datastore.GetDB(self.ConfigObj)
-	require.NoError(t, err)
+	closer := utils.SetFlowIdForTests("F.SomeFlow")
+	defer closer()
+
+	// Create a new collection flow
+	flow_id, err := self.createArtifactCollection()
+	require.NoError(self.T(), err)
+	assert.Equal(self.T(), "F.SomeFlow", flow_id)
+
+	launcher, err := services.GetLauncher(self.ConfigObj)
+	assert.NoError(t, err)
+
+	collection_context, err := launcher.GetFlowDetails(self.Ctx, self.ConfigObj,
+		services.GetFlowOptions{}, self.client_id, flow_id)
+	assert.NoError(t, err)
+
+	assert.Equal(self.T(), collection_context.Context.State,
+		flows_proto.ArtifactCollectorContext_RUNNING)
 
 	runner := flows.NewFlowRunner(self.Ctx, self.ConfigObj)
 	defer runner.Close(self.Ctx)
 
 	// Send a message to a random non-existant flow from client.
-	flow_id := "F.NONEXISTENT"
 	runner.ProcessSingleMessage(
 		self.Ctx,
 		&crypto_proto.VeloMessage{
 			Source:    self.client_id,
 			SessionId: flow_id,
 			FlowStats: &crypto_proto.FlowStats{
 				QueryStatus: []*crypto_proto.VeloStatus{
-					{Status: crypto_proto.VeloStatus_OK, QueryId: 1},
+					{
+						Status:       crypto_proto.VeloStatus_UNKNOWN_FLOW,
+						ErrorMessage: "Unknown flow",
+					},
 				},
 			},
 		})
 
-	// We used to send cancellation message to the client, but this
-	// too expensive for the server to keep track of. Now we just
-	// write data in the flow as if it exists anyway.
+	collection_context, err = launcher.GetFlowDetails(self.Ctx, self.ConfigObj,
+		services.GetFlowOptions{}, self.client_id, flow_id)
+	assert.NoError(t, err)
 
-	// The flow does not exist - make sure it still does not.
-	collection_context := &flows_proto.ArtifactCollectorContext{}
-	path_manager := paths.NewFlowPathManager(self.client_id, flow_id)
-	err = db.GetSubject(self.ConfigObj, path_manager.Path(), collection_context)
-	require.Error(t, err, os.ErrNotExist)
+	assert.Equal(self.T(), collection_context.Context.State,
+		flows_proto.ArtifactCollectorContext_ERROR)
 
-	// The flow stats are written as normal.
-	err = db.GetSubject(self.ConfigObj, path_manager.Stats(), collection_context)
-	assert.NoError(t, err)
+	assert.Equal(self.T(), collection_context.Context.Status,
+		"Unknown flow")
+
+	assert.Equal(self.T(), len(collection_context.Context.QueryStats), 1)
 }
 
 // Test an unknown flow. What happens when the server receives a

services/client_info/client_info_test.go

@@ -48,6 +48,11 @@ type: INTERNAL
 `, `
 name: Server.Audit.Logs
 type: INTERNAL
+`, `
+name: Client.Test
+type: CLIENT
+sources:
+- query: SELECT * FROM info()
 `})
 
 	// Create a client in the datastore so we can test initializing

services/client_info/fixtures/TestInFlightMessages.golden

@@ -0,0 +1,42 @@
+{
+ "InFlightFlows": {
+  "client_id": "C.1234",
+  "hostname": "Hostname",
+  "has_tasks": true,
+  "in_flight_flows": {
+   "F.0": 10,
+   "F.1": 10,
+   "F.2": 10,
+   "F.3": 10
+  }
+ },
+ "StatusChecks": [
+  {
+   "session_id": "F.Status",
+   "flow_stats_request": {
+    "flow_id": [
+     "F.0",
+     "F.1",
+     "F.2",
+     "F.3"
+    ]
+   }
+  }
+ ],
+ "AfterCompletion": {
+  "client_id": "C.1234",
+  "hostname": "Hostname",
+  "has_tasks": true
+ },
+ "SecondSetOfTasks": {
+  "client_id": "C.1234",
+  "hostname": "Hostname",
+  "has_tasks": true,
+  "in_flight_flows": {
+   "F.4": 100,
+   "F.5": 100,
+   "F.6": 100,
+   "F.7": 100
+  }
+ }
+}

services/client_info/tasks_test.go

@@ -6,13 +6,19 @@ import (
 	"sort"
 	"time"
 
+	"github.com/Velocidex/ordereddict"
 	"google.golang.org/protobuf/proto"
 	crypto_proto "www.velocidex.com/golang/velociraptor/crypto/proto"
+	"www.velocidex.com/golang/velociraptor/flows"
+	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
+	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/services"
 	"www.velocidex.com/golang/velociraptor/services/client_info"
 	"www.velocidex.com/golang/velociraptor/utils"
+	"www.velocidex.com/golang/velociraptor/vql/acl_managers"
 	"www.velocidex.com/golang/velociraptor/vtesting"
 	"www.velocidex.com/golang/velociraptor/vtesting/assert"
+	"www.velocidex.com/golang/velociraptor/vtesting/goldie"
 )
 
 func (self *ClientInfoTestSuite) TestQueueMessages() {
@@ -85,3 +91,111 @@ func (self *ClientInfoTestSuite) TestFastQueueMessages() {
 		assert.True(self.T(), proto.Equal(tasks[i], written[i]))
 	}
 }
+
+func (self *ClientInfoTestSuite) TestInFlightMessages() {
+	closer := utils.MockTime(utils.NewMockClock(time.Unix(10, 0)))
+	defer closer()
+
+	launcher, err := services.GetLauncher(self.ConfigObj)
+	assert.NoError(self.T(), err)
+
+	var flow_ids []string
+	acl_manager := acl_managers.NullACLManager{}
+	manager, _ := services.GetRepositoryManager(self.ConfigObj)
+	repository, _ := manager.GetGlobalRepository(self.ConfigObj)
+
+	for i := 0; i < 10; i++ {
+		closer := utils.SetFlowIdForTests(fmt.Sprintf("F.%d", i))
+
+		flow_id, err := launcher.ScheduleArtifactCollection(self.Ctx,
+			self.ConfigObj, acl_manager,
+			repository, &flows_proto.ArtifactCollectorArgs{
+				Creator:   "admin",
+				ClientId:  self.client_id,
+				Artifacts: []string{"Client.Test"},
+			}, utils.SyncCompleter)
+		assert.NoError(self.T(), err)
+
+		flow_ids = append(flow_ids, flow_id)
+
+		closer()
+	}
+
+	client_info_manager, err := services.GetClientInfoManager(self.ConfigObj)
+	assert.NoError(self.T(), err)
+
+	tasks, err := client_info_manager.GetClientTasks(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	// 4 tasks are queued
+	assert.Equal(self.T(), len(tasks), 4)
+
+	tasks, err = client_info_manager.GetClientTasks(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	// Tasks are still in flight, so we can not get any new tasks yet.
+	assert.Equal(self.T(), len(tasks), 0)
+
+	client_info, err := client_info_manager.Get(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	// Should contain only 4 flow ids in the in_flight_flows set.
+	golden := ordereddict.NewDict().
+		Set("InFlightFlows", client_info)
+
+	// Pass some time
+	closer = utils.MockTime(utils.NewMockClock(time.Unix(100, 0)))
+	defer closer()
+
+	// Tasks are still in flight, so we do not send any flows, instead
+	// we send a task status request to see how those other tasks are
+	// going.
+	tasks, err = client_info_manager.GetClientTasks(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	sort.Strings(tasks[0].FlowStatsRequest.FlowId)
+
+	assert.Equal(self.T(), len(tasks), 1)
+
+	// Should contains a status check request for all inflight flows.
+	golden.Set("StatusChecks", tasks)
+
+	// Now complete the flows
+	runner := flows.NewFlowRunner(self.Ctx, self.ConfigObj)
+	for flow_id := range client_info.InFlightFlows {
+		runner.ProcessSingleMessage(self.Ctx,
+			&crypto_proto.VeloMessage{
+				Source:    self.client_id,
+				SessionId: flow_id,
+				FlowStats: &crypto_proto.FlowStats{
+					FlowComplete: true,
+					QueryStatus: []*crypto_proto.VeloStatus{{
+						Status: crypto_proto.VeloStatus_OK,
+					}},
+				}})
+	}
+	runner.Close(self.Ctx)
+
+	client_info, err = client_info_manager.Get(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	// Completing the flows removes the flows from the in flight set.
+	assert.Equal(self.T(), len(client_info.InFlightFlows), 0)
+
+	// Should contain no in flight flows but still contain the
+	// has_tasks flag.
+	golden.Set("AfterCompletion", client_info)
+
+	// Now read some more tasks
+	tasks, err = client_info_manager.GetClientTasks(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+
+	// Should conatin
+	assert.Equal(self.T(), len(tasks), 4)
+
+	client_info, err = client_info_manager.Get(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+	golden.Set("SecondSetOfTasks", client_info)
+
+	goldie.Assert(self.T(), "TestInFlightMessages", json.MustMarshalIndent(golden))
+}

services/hunt_dispatcher/storage.go

@@ -561,7 +561,7 @@ func (self *HuntStorageManagerImpl) loadHuntsFromDatastore(
 			self.dirty = true
 
 			// The old hunt record is newer than the one on disk, ignore it.
-		} else if old_hunt_record.Version >= hunt_obj.Version {
+		} else if old_hunt_record.Version > hunt_obj.Version {
 			continue
 		}
 

services/launcher/flows.go

@@ -276,6 +276,15 @@ func UpdateFlowStats(collection_context *flows_proto.ArtifactCollectorContext) {
 			collection_context.StartTime = s.FirstActive
 		}
 
+		// If the Query stats represents an unknown flow, we mark the
+		// flow as errored.
+		if s.Status == crypto_proto.VeloStatus_UNKNOWN_FLOW {
+			collection_context.State = flows_proto.ArtifactCollectorContext_ERROR
+			collection_context.Status = s.ErrorMessage
+			collection_context.Backtrace = s.Backtrace
+			break
+		}
+
 		// Get the first errored query and mark the entire collection_context with it.
 		if collection_context.State == flows_proto.ArtifactCollectorContext_RUNNING &&
 			s.Status == crypto_proto.VeloStatus_GENERIC_ERROR {

vql/functions/alerts.go

@@ -46,7 +46,7 @@ func (self *AlertFunction) Call(ctx context.Context,
 
 	alert_name, pres := args.GetString("name")
 	if !pres {
-		scope.Log("alert: Alert name must be specified!")
+		scope.Log("ERROR:alert: Alert name must be specified!")
 		return &vfilter.Null{}
 	}