Add Modify() method to client info manager. (#2933)

This is required to ensure client info records can be modified
atomically. Fixes certain edge cases arising due to server monitoring
queries assigning labels racing with interrogations.

.github/workflows/go.yml

@@ -12,7 +12,7 @@ jobs:
     name: Build
     # We used to build on ubuntu-18.04 but that is now deprecated by
     # GitHub. Earlier distributions will have to use the musl build.
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04
     steps:
 
     - name: Check out code into the Go module directory

artifacts/definitions/Server/Utils/CreateCollector.yaml

@@ -349,7 +349,9 @@ parameters:
                 remove_last=TRUE,
                 permissions=if(condition=IsExecutable, then="x")))
 
-       SELECT copy(filename=Filename, accessor="me", dest=temp_binary) AS OSPath,
+       LET FullPath <= copy(filename=Filename, accessor="me", dest=temp_binary)
+
+       SELECT FullPath, FullPath AS OSPath,
               Filename AS Name
        FROM matching_tools
 

artifacts/proto/artifact.proto

@@ -202,7 +202,7 @@ message Artifact {
     // Artifact was already compiled.
     bool compiled = 14;
 
-    // The artifact is considered built and and can not be overriden by the GUI
+    // The artifact is considered built in and can not be overriden by the GUI
     bool built_in = 20;
 
     // The artifact is built into the Velociraptor binary.

services/client_info.go

@@ -64,10 +64,18 @@ func (self ClientInfo) OS() ClientOS {
 type ClientInfoManager interface {
 	ListClients(ctx context.Context) <-chan string
 
-	// Used to set a new client record.
+	// Used to set a new client record. To modify an existing record -
+	// or set a new one use Modify()
 	Set(ctx context.Context,
 		client_info *ClientInfo) error
 
+	// Modify a record or set a new one - if the record is not found,
+	// modifier will receive a nil client_info. The ClientInfoManager
+	// can not be accessed within the modifier function as it is
+	// locked for the duration of the change.
+	Modify(ctx context.Context, client_id string,
+		modifier func(client_info *ClientInfo) (new_record *ClientInfo, err error)) error
+
 	Get(ctx context.Context,
 		client_id string) (*ClientInfo, error)
 

services/client_info/client_info.go

@@ -414,6 +414,13 @@ func (self *ClientInfoManager) ProcessPing(
 	return nil
 }
 
+func (self *ClientInfoManager) Modify(
+	ctx context.Context, client_id string,
+	modifier func(client_info *services.ClientInfo) (
+		*services.ClientInfo, error)) error {
+	return self.storage.Modify(ctx, client_id, modifier)
+}
+
 func (self *ClientInfoManager) Get(
 	ctx context.Context, client_id string) (*services.ClientInfo, error) {
 	record, err := self.storage.GetRecord(client_id)

services/client_info/client_info_test.go

@@ -56,6 +56,48 @@ type: INTERNAL
 	self.TestSuite.SetupTest()
 }
 
+func (self *ClientInfoTestSuite) TestClientInfoModify() {
+	// Fetch the client from the manager
+	client_info_manager, err := services.GetClientInfoManager(self.ConfigObj)
+	assert.NoError(self.T(), err)
+
+	info, err := client_info_manager.Get(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+	assert.Equal(self.T(), info.ClientId, self.client_id)
+	assert.Equal(self.T(), info.Ping, uint64(0))
+
+	// Update the ping time
+	err = client_info_manager.Modify(self.Ctx, self.client_id,
+		func(client_info *services.ClientInfo) (*services.ClientInfo, error) {
+			assert.NotNil(self.T(), client_info)
+
+			client_info.Ping = 10
+			return client_info, nil
+		})
+	assert.NoError(self.T(), err)
+
+	// Now get the client record and check that it is updated
+	info, err = client_info_manager.Get(self.Ctx, self.client_id)
+	assert.NoError(self.T(), err)
+	assert.Equal(self.T(), info.Ping, uint64(10))
+
+	// Now modify a nonexistant client - equivalent to a Set() call
+	// (atomic check and set).
+	err = client_info_manager.Modify(self.Ctx, "C.DOESNOTEXIT",
+		func(client_info *services.ClientInfo) (*services.ClientInfo, error) {
+			assert.Nil(self.T(), client_info)
+			return &services.ClientInfo{actions_proto.ClientInfo{
+				ClientId: "C.DOESNOTEXIT",
+				Ping:     20,
+			}}, nil
+		})
+	assert.NoError(self.T(), err)
+
+	info, err = client_info_manager.Get(self.Ctx, "C.DOESNOTEXIT")
+	assert.NoError(self.T(), err)
+	assert.Equal(self.T(), info.Ping, uint64(20))
+}
+
 func (self *ClientInfoTestSuite) TestClientInfo() {
 	// Fetch the client from the manager
 	client_info_manager, err := services.GetClientInfoManager(self.ConfigObj)

services/client_info/storage.go

@@ -79,10 +79,45 @@ func (self *Store) Remove(client_id string) {
 	delete(self.data, client_id)
 }
 
+func (self *Store) Modify(
+	ctx context.Context, client_id string,
+	modifier func(client_info *services.ClientInfo) (
+		*services.ClientInfo, error)) error {
+
+	self.mu.Lock()
+	defer self.mu.Unlock()
+
+	client_info, _ := self._GetRecord(client_id)
+	var record *services.ClientInfo
+
+	if client_info != nil {
+		record = &services.ClientInfo{*client_info}
+	}
+
+	// If the record was not changed just ignore it.
+	new_record, err := modifier(record)
+	if err != nil {
+		return err
+	}
+
+	// Callback can indicate no change is needed by returning a nil
+	// for client_info.
+	if new_record == nil {
+		return nil
+	}
+
+	// Write the modified record to the LRU
+	return self._SetRecord(&new_record.ClientInfo)
+}
+
 func (self *Store) GetRecord(client_id string) (*actions_proto.ClientInfo, error) {
 	self.mu.Lock()
 	defer self.mu.Unlock()
 
+	return self._GetRecord(client_id)
+}
+
+func (self *Store) _GetRecord(client_id string) (*actions_proto.ClientInfo, error) {
 	serialized, pres := self.data[client_id]
 	if !pres {
 		return nil, os.ErrNotExist
@@ -106,6 +141,10 @@ func (self *Store) SetRecord(record *actions_proto.ClientInfo) error {
 	self.mu.Lock()
 	defer self.mu.Unlock()
 
+	return self._SetRecord(record)
+}
+
+func (self *Store) _SetRecord(record *actions_proto.ClientInfo) error {
 	serialized, err := proto.Marshal(record)
 	if err != nil {
 		return err