Bugfix: In flight flow calculation was stuck sometimes

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{}
 	}