Maintain row count per query. (Velocidex#2113)

This fixes a bug where row count was not consistent for aritfacts with
multiple sources because the row count was maintained for each request
but a request may contain multiple queries.

This only affects CloudVelo

actions/tracker.go

@@ -0,0 +1,36 @@
+package actions
+
+import (
+	"sync"
+
+	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
+)
+
+type QueryTracker struct {
+	mu sync.Mutex
+
+	queriesToStartRow map[string]uint64
+}
+
+func (self *QueryTracker) GetStartRow(query *actions_proto.VQLRequest) uint64 {
+	self.mu.Lock()
+	defer self.mu.Unlock()
+
+	start_row, _ := self.queriesToStartRow[query.Name]
+	return start_row
+}
+
+func (self *QueryTracker) AddRows(
+	query *actions_proto.VQLRequest, count uint64) {
+	self.mu.Lock()
+	defer self.mu.Unlock()
+
+	start_row, _ := self.queriesToStartRow[query.Name]
+	self.queriesToStartRow[query.Name] = start_row + count
+}
+
+func NewQueryTracker() *QueryTracker {
+	return &QueryTracker{
+		queriesToStartRow: make(map[string]uint64),
+	}
+}

actions/vql.go

@@ -196,7 +196,7 @@ func (self VQLClientAction) StartQuery(
 		return
 	}
 
-	start_row := 0
+	row_tracker := NewQueryTracker()
 
 	// All the queries will use the same scope. This allows one
 	// query to define functions for the next query in order.
@@ -258,11 +258,11 @@ func (self VQLClientAction) StartQuery(
 					Part:          uint64(result.Part),
 					JSONLResponse: string(result.Payload),
 					TotalRows:     uint64(result.TotalRows),
-					QueryStartRow: uint64(start_row),
+					QueryStartRow: row_tracker.GetStartRow(query),
 					Timestamp:     uint64(time.Now().UTC().UnixNano() / 1000),
 				}
 
-				start_row += result.TotalRows
+				row_tracker.AddRows(query, uint64(result.TotalRows))
 
 				// Don't log empty VQL statements.
 				if query.Name != "" {

api/api.go

@@ -52,6 +52,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/file_store/path_specs"
 	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
 	"www.velocidex.com/golang/velociraptor/grpc_client"
+	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/paths"
 	"www.velocidex.com/golang/velociraptor/server"
@@ -111,7 +112,7 @@ func (self *ApiServer) CancelFlow(
 			"user":    user_name,
 			"client":  in.ClientId,
 			"flow_id": in.FlowId,
-			"details": fmt.Sprintf("%v", in),
+			"details": json.MustMarshalString(in),
 		}).Info("CancelFlow")
 
 	return result, nil
@@ -216,7 +217,7 @@ func (self *ApiServer) CollectArtifact(
 			"user":    in.Creator,
 			"client":  in.ClientId,
 			"flow_id": flow_id,
-			"details": fmt.Sprintf("%v", in),
+			"details": json.MustMarshalString(in),
 		}).Info("CollectArtifact")
 
 	return result, nil

api/hunts.go

@@ -1,7 +1,6 @@
 package api
 
 import (
-	"fmt"
 	"time"
 
 	"github.com/Velocidex/ordereddict"
@@ -14,6 +13,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/acls"
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	"www.velocidex.com/golang/velociraptor/file_store/csv"
+	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
 	"www.velocidex.com/golang/velociraptor/services/hunt_dispatcher"
@@ -116,7 +116,7 @@ func (self *ApiServer) CreateHunt(
 		WithFields(logrus.Fields{
 			"user":    in.Creator,
 			"hunt_id": in.HuntId,
-			"details": fmt.Sprintf("%v", in),
+			"details": json.MustMarshalString(in),
 		}).Info("CreateHunt")
 
 	result := &api_proto.StartFlowResponse{}
@@ -161,7 +161,7 @@ func (self *ApiServer) ModifyHunt(
 		WithFields(logrus.Fields{
 			"user":    in.Creator,
 			"hunt_id": in.HuntId,
-			"details": fmt.Sprintf("%v", in),
+			"details": json.MustMarshalString(in),
 		}).Info("ModifyHunt")
 
 	hunt_dispatcher, err := services.GetHuntDispatcher(org_config_obj)