Added tests for ring buffer. (Velocidex#216)

artifacts/artifacts.go

@@ -79,7 +79,7 @@ func (self *Repository) LoadDirectory(dirname string) (*int, error) {
 		})
 }
 
-var query_regexp = regexp.MustCompile(`(?im)(^ +- +)(SELECT|LET)`)
+var query_regexp = regexp.MustCompile(`(?im)(^ +- +)(SELECT|LET|//)`)
 
 // Fix common YAML errors.
 func sanitize_artifact_yaml(data string) string {

artifacts/definitions/Admin/Client/Upgrade.yaml

@@ -17,8 +17,12 @@ sources:
   - precondition:
       SELECT OS From info() where OS = 'windows'
     queries:
-      - // Run the installer on the downloaded file.
-        SELECT * from foreach(
+      # Wait a random amount of time so this can be run in a
+      # hunt. Otherwise all clients will attempt to download the same
+      # file at the same time probably overloading the server.
+      - LET _ = <= SELECT sleep(time=rand(range=600)) FROM scope()
+
+      - SELECT * from foreach(
          row={
             SELECT Content AS Binary
             FROM http_client(url=clientURL, tempfile_extension=".msi")

config/config.go

@@ -75,7 +75,7 @@ func GetDefaultConfig() *config_proto.Config {
 			LocalBuffer: &config_proto.RingBufferConfig{
 				MemorySize: 50 * 1024 * 1024,
 				DiskSize:   1024 * 1024 * 1024,
-				Filename:   "$Temp/Velociraptor_Buffer.bin",
+				Filename:   "/tmp/Velociraptor_Buffer.bin",
 			},
 
 			// Specific instructions for the
@@ -154,6 +154,7 @@ func GetDefaultConfig() *config_proto.Config {
 	if runtime.GOOS == "windows" {
 		result.Datastore.Location = "C:\\Windows\\Temp"
 		result.Datastore.FilestoreDirectory = "C:\\Windows\\Temp"
+		result.Client.LocalBuffer.Filename = "C:\\Windows\\Temp\\Velociraptor_Buffer.bin"
 	}
 
 	return result

datastore/filebased.go

@@ -37,6 +37,7 @@
 package datastore
 
 import (
+	"compress/gzip"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -527,17 +528,33 @@ func readContentFromFile(
 	}
 
 	file, err := os.Open(filename)
-	if err != nil {
-		if !must_exist && os.IsNotExist(err) {
-			return []byte{}, nil
+	if err == nil {
+		defer file.Close()
+
+		result, err := ioutil.ReadAll(
+			io.LimitReader(file, constants.MAX_MEMORY))
+		return result, errors.WithStack(err)
+	}
+
+	// File does not exist - try the gzip version
+	if os.IsNotExist(err) {
+		file, err = os.Open(filename + ".gz")
+		if err == nil {
+			zr, err := gzip.NewReader(file)
+			if err != nil {
+				return nil, errors.WithStack(err)
+			}
+			result, err := ioutil.ReadAll(
+				io.LimitReader(zr, constants.MAX_MEMORY))
+			return result, errors.WithStack(err)
 		}
-		return nil, errors.WithStack(err)
 	}
-	defer file.Close()
 
-	result, err := ioutil.ReadAll(
-		io.LimitReader(file, constants.MAX_MEMORY))
-	return result, errors.WithStack(err)
+	// Its ok if the file does not exist - no error.
+	if !must_exist && os.IsNotExist(err) {
+		return []byte{}, nil
+	}
+	return nil, errors.WithStack(err)
 }
 
 // Convert a file name from the data store to a urn.

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/Velocidex/ordereddict v0.0.0-20191106020901-97c468e5e403
 	github.com/Velocidex/survey v1.8.7-0.20190926071832-2ff99cc7aa49
 	github.com/Velocidex/yaml v0.0.0-20190812045153-ad0acda9eea0
+	github.com/alecthomas/assert v0.0.0-20170929043011-405dbfeb8e38
 	github.com/alecthomas/chroma v0.6.0
 	github.com/alecthomas/participle v0.4.1
 	github.com/alexmullins/zip v0.0.0-20180717182244-4affb64b04d0