Bugfix: Fixed crash with timeout handling on API handler (#3187)

Race condition could result in logging messages being sent after the
connection is already closed.

.github/workflows/windows.yml

@@ -1,5 +1,6 @@
 name: Windows Test
 on: [pull_request]
+
 jobs:
   build:
     name: Windows Test
@@ -16,7 +17,7 @@ jobs:
 
     - name: Configure test environment
       shell: cmd
-      if: success()
+      if: always()
       run: |
         echo %PATH%
         echo %GOPATH%
@@ -36,7 +37,7 @@ jobs:
         regedit /S artifacts/testdata/windows/init.reg
 
     - name: Build
-      if: success()
+      if: always()
       env:
         CC: x86_64-w64-mingw32-gcc
       shell: bash
@@ -65,7 +66,7 @@ jobs:
 
     - name: Test
       shell: bash
-      if: success()
+      if: always()
       env:
         # Disable CGO for building tests - it takes too long and it is
         # not needed (mainly disables Yara building again).

api/query.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"log"
 	"runtime/debug"
+	"sync"
 	"time"
 
 	"github.com/Velocidex/ordereddict"
@@ -36,6 +37,7 @@ import (
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
+	"www.velocidex.com/golang/velociraptor/utils"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
 	"www.velocidex.com/golang/velociraptor/vql/acl_managers"
 	"www.velocidex.com/golang/vfilter"
@@ -99,6 +101,10 @@ func streamQuery(
 	// Now execute the query.
 	scope := manager.BuildScope(builder)
 
+	// Keep the connection open until the logs are sent.
+	wg := sync.WaitGroup{}
+	defer wg.Wait()
+
 	subctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -108,13 +114,19 @@ func streamQuery(
 		timed_ctx, timed_cancel := context.WithTimeout(subctx,
 			time.Second*time.Duration(arg.Timeout))
 
+		wg.Add(1)
 		go func() {
+			defer wg.Done()
+
 			select {
-			case <-ctx.Done():
+			// Cancelling the parent will not return a log.
+			case <-subctx.Done():
 				timed_cancel()
+
+				// Log the timeout
 			case <-timed_ctx.Done():
 				scope.Log("collect: <red>Timeout Error:</> Collection timed out after %v",
-					time.Now().Sub(start))
+					utils.GetTime().Now().Sub(start))
 				// Cancel the main context.
 				cancel()
 				timed_cancel()
@@ -126,7 +138,9 @@ func streamQuery(
 	scope.SetThrottler(
 		actions.NewThrottler(subctx, scope, 0, float64(arg.CpuLimit), 0))
 
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer close(response_channel)
 		defer scope.Close()
 

file_store/directory/listener_test.go

@@ -73,7 +73,7 @@ func (self *TestSuite) TestListenerPreserveTypes() {
 		Set("B", uint64(9223372036854775808))
 	listener.Send(event_source)
 
-	vtesting.WaitUntil(time.Second, self.T(), func() bool {
+	vtesting.WaitUntil(time.Second*5, self.T(), func() bool {
 		// Make sure we wrote to the buffer file.
 		size, _ := listener.Debug().GetInt64("Size")
 		return size > directory.FirstRecordOffset

services/interrogation/interrogation_test.go

@@ -2,6 +2,7 @@ package interrogation_test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -15,6 +16,7 @@ import (
 	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
 	"www.velocidex.com/golang/velociraptor/paths"
 	"www.velocidex.com/golang/velociraptor/services"
+	"www.velocidex.com/golang/velociraptor/utils"
 	"www.velocidex.com/golang/velociraptor/vtesting"
 
 	_ "www.velocidex.com/golang/velociraptor/result_sets/timed"
@@ -37,7 +39,7 @@ type: INTERNAL
 `,
 	})
 
-	self.client_id = "C.12312"
+	self.client_id = fmt.Sprintf("C.1%d", utils.GetId())
 	self.flow_id = "F.1232"
 
 	self.TestSuite.SetupTest()
@@ -180,6 +182,10 @@ func (self *ServicesTestSuite) TestEnrollService() {
 		}
 	}
 
+	if len(children) > 1 {
+		test_utils.GetMemoryDataStore(self.T(), self.ConfigObj).Debug(self.ConfigObj)
+	}
+
 	assert.Equal(self.T(), len(children), 1)
 	assert.Equal(self.T(), children[0].Base(),
 		client_info.LastInterrogateFlowId)