Fixed bug in uploading sparse files. (Velocidex#519)

* Fixed bug in uploading sparse files.

This affected files fetched with NTFS which spanned multiple runs.

* Added test.

gui/static/angular-components/client/shell-viewer-directive.js

@@ -44,7 +44,7 @@ ShellViewerController.prototype.launchCommand = function() {
         artifact = "Windows.System.PowerShell";
     } else if(this.type == "Cmd") {
         artifact = "Windows.System.CmdShell";
-    } else if(this.type == "Base") {
+    } else if(this.type == "Bash") {
         artifact = "Linux.Sys.BashShell";
     } else {
         return;
@@ -91,7 +91,8 @@ ShellViewerController.prototype.fetchLastShellCollections = function() {
                 for (var j=0; j<artifacts.length; j++) {
                     var artifact = artifacts[j];
                     if (artifact == "Windows.System.PowerShell" ||
-                        artifact == "Windows.System.CmdShell") {
+                        artifact == "Windows.System.CmdShell" ||
+                        artifact == "Linux.Sys.BashShell" ) {
                         self.flows.push(items[i]);
                     }
                 }

json/wrappers.go

@@ -19,6 +19,14 @@ func Marshal(v interface{}) ([]byte, error) {
 	return json.MarshalWithOptions(v, opts)
 }
 
+func MustMarshalIndent(v interface{}) []byte {
+	result, err := MarshalIndent(v)
+	if err != nil {
+		panic(err)
+	}
+	return result
+}
+
 func MarshalIndent(v interface{}) ([]byte, error) {
 	opts := NewEncOpts()
 	return MarshalIndentWithOptions(v, opts)

responder/responder.go

@@ -32,7 +32,7 @@ import (
 )
 
 type Responder struct {
-	output chan<- *crypto_proto.GrrMessage
+	output chan *crypto_proto.GrrMessage
 
 	sync.Mutex
 	request *crypto_proto.GrrMessage
@@ -44,7 +44,7 @@ type Responder struct {
 func NewResponder(
 	config_obj *config_proto.Config,
 	request *crypto_proto.GrrMessage,
-	output chan<- *crypto_proto.GrrMessage) *Responder {
+	output chan *crypto_proto.GrrMessage) *Responder {
 	result := &Responder{
 		request: request,
 		next_id: 0,

responder/testutils.go

@@ -0,0 +1,20 @@
+package responder
+
+import crypto_proto "www.velocidex.com/golang/velociraptor/crypto/proto"
+
+func TestResponder() *Responder {
+	return &Responder{
+		output:  make(chan *crypto_proto.GrrMessage, 100),
+		request: &crypto_proto.GrrMessage{},
+	}
+}
+
+func GetTestResponses(self *Responder) []*crypto_proto.GrrMessage {
+	close(self.output)
+	result := []*crypto_proto.GrrMessage{}
+	for item := range self.output {
+		result = append(result, item)
+	}
+
+	return result
+}

uploads/api.go

@@ -11,8 +11,9 @@ import (
 
 type Range struct {
 	// In bytes
-	Offset, Length int64
-	IsSparse       bool
+	Offset   int64
+	Length   int64
+	IsSparse bool
 }
 
 type RangeReader interface {

uploads/client_uploader.go

@@ -19,8 +19,8 @@ import (
 	"www.velocidex.com/golang/vfilter"
 )
 
-const (
-	BUFF_SIZE = 1024 * 1024
+var (
+	BUFF_SIZE = int64(1024 * 1024)
 )
 
 // An uploader delivering files from client to server.
@@ -198,26 +198,26 @@ func (self *VelociraptorUploader) maybeUploadSparse(
 		range_reader.Seek(read_offset, os.SEEK_SET)
 
 		for to_read > 0 {
+			to_read_buf := to_read
+
 			// Ensure there is a fresh allocation for every
 			// iteration to prevent overwriting in-flight buffers.
-			if to_read > BUFF_SIZE {
-				to_read = BUFF_SIZE
+			if to_read_buf > BUFF_SIZE {
+				to_read_buf = BUFF_SIZE
 			}
 
-			buffer := make([]byte, BUFF_SIZE)
+			buffer := make([]byte, to_read_buf)
 			read_bytes, err := range_reader.Read(buffer)
 			// Hard read error - give up.
 			if err != nil && err != io.EOF {
 				return nil, err
 			}
-
 			// End of range - go to the next range
 			if read_bytes == 0 {
-				break
+				continue
 			}
 
 			data := buffer[:read_bytes]
-
 			sha_sum.Write(data)
 			md5_sum.Write(data)
 
@@ -246,7 +246,6 @@ func (self *VelociraptorUploader) maybeUploadSparse(
 			to_read -= int64(read_bytes)
 			write_offset += int64(read_bytes)
 			read_offset += int64(read_bytes)
-
 		}
 	}
 
@@ -256,7 +255,6 @@ func (self *VelociraptorUploader) maybeUploadSparse(
 		if err != nil {
 			return nil, err
 		}
-
 		self.Responder.AddResponse(&crypto_proto.GrrMessage{
 			RequestId: constants.TransferWellKnownFlowId,
 			FileBuffer: &actions_proto.FileBuffer{

uploads/client_uploader_test.go

@@ -0,0 +1,131 @@
+package uploads
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/alecthomas/assert"
+	"github.com/sebdah/goldie"
+	crypto_proto "www.velocidex.com/golang/velociraptor/crypto/proto"
+	"www.velocidex.com/golang/velociraptor/json"
+	"www.velocidex.com/golang/velociraptor/responder"
+	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+)
+
+type TestRangeReader struct {
+	*bytes.Reader
+	ranges []Range
+}
+
+func (self *TestRangeReader) Ranges() []Range {
+	return self.ranges
+}
+
+// Combine the output of all fragments into a strings
+func CombineOutput(name string, responses []*crypto_proto.GrrMessage) string {
+	result := []byte{}
+
+	for _, item := range responses {
+		if item.FileBuffer.Pathspec.Path == name {
+			result = append(result, item.FileBuffer.Data...)
+		}
+	}
+
+	return string(result)
+}
+
+func TestClientUploaderSparse(t *testing.T) {
+	resp := responder.TestResponder()
+	uploader := &VelociraptorUploader{
+		Responder: resp,
+	}
+
+	BUFF_SIZE = 10000
+
+	reader := &TestRangeReader{
+		Reader: bytes.NewReader([]byte(
+			"Hello world hello world")),
+		ranges: []Range{
+			{Offset: 0, Length: 6, IsSparse: false},
+			{Offset: 6, Length: 6, IsSparse: true},
+			{Offset: 12, Length: 6, IsSparse: false},
+		},
+	}
+	range_reader, ok := interface{}(reader).(RangeReader)
+	assert.Equal(t, ok, true)
+	ctx := context.Background()
+	scope := vql_subsystem.MakeScope()
+	uploader.maybeUploadSparse(ctx, scope,
+		"foo", "ntfs", "", 1000, range_reader)
+	responses := responder.GetTestResponses(resp)
+
+	// Expected size is the combined sum of all ranges with data
+	// in them
+	assert.Equal(t, responses[0].FileBuffer.Size, uint64(12))
+
+	assert.Equal(t, CombineOutput("foo", responses),
+		"Hello hello ")
+	goldie.Assert(t, "ClientUploaderSparse",
+		json.MustMarshalIndent(responses))
+	assert.NotEqual(t, CombineOutput("foo.idx", responses), "")
+}
+
+func TestClientUploaderSparseMultiBuffer(t *testing.T) {
+	resp := responder.TestResponder()
+	uploader := &VelociraptorUploader{
+		Responder: resp,
+	}
+
+	// 2 bytes per message
+	BUFF_SIZE = 2
+	reader := &TestRangeReader{
+		Reader: bytes.NewReader([]byte(
+			"Hello world hello world")),
+		ranges: []Range{
+			{Offset: 0, Length: 6, IsSparse: false},
+			{Offset: 6, Length: 6, IsSparse: true},
+			{Offset: 12, Length: 6, IsSparse: false},
+		},
+	}
+	range_reader, ok := interface{}(reader).(RangeReader)
+	assert.Equal(t, ok, true)
+	ctx := context.Background()
+	scope := vql_subsystem.MakeScope()
+	uploader.maybeUploadSparse(ctx, scope,
+		"foo", "ntfs", "", 1000, range_reader)
+	responses := responder.GetTestResponses(resp)
+	assert.Equal(t, CombineOutput("foo", responses), "Hello hello ")
+	goldie.Assert(t, "ClientUploaderSparseMultiBuffer",
+		json.MustMarshalIndent(responses))
+	assert.NotEqual(t, CombineOutput("foo.idx", responses), "")
+}
+
+func TestClientUploaderNoIndexIfNotSparse(t *testing.T) {
+	resp := responder.TestResponder()
+	uploader := &VelociraptorUploader{
+		Responder: resp,
+	}
+
+	// 2 bytes per message
+	BUFF_SIZE = 2
+	reader := &TestRangeReader{
+		Reader: bytes.NewReader([]byte(
+			"Hello world hello world")),
+		ranges: []Range{
+			{Offset: 0, Length: 6, IsSparse: false},
+			{Offset: 12, Length: 6, IsSparse: false},
+		},
+	}
+	range_reader, ok := interface{}(reader).(RangeReader)
+	assert.Equal(t, ok, true)
+	ctx := context.Background()
+	scope := vql_subsystem.MakeScope()
+	uploader.maybeUploadSparse(ctx, scope,
+		"foo", "ntfs", "", 1000, range_reader)
+	responses := responder.GetTestResponses(resp)
+	assert.Equal(t, CombineOutput("foo", responses), "Hello hello ")
+
+	// No idx written when there are no sparse ranges.
+	assert.Equal(t, CombineOutput("foo.idx", responses), "")
+}

uploads/fixtures/ClientUploaderSparse.golden

@@ -0,0 +1,40 @@
+[
+ {
+  "request_id": 5,
+  "FileBuffer": {
+   "pathspec": {
+    "path": "foo",
+    "accessor": "ntfs"
+   },
+   "size": 12,
+   "data": "SGVsbG8g"
+  }
+ },
+ {
+  "request_id": 5,
+  "response_id": 1,
+  "FileBuffer": {
+   "pathspec": {
+    "path": "foo",
+    "accessor": "ntfs"
+   },
+   "offset": 6,
+   "size": 12,
+   "data": "aGVsbG8g",
+   "eof": true
+  }
+ },
+ {
+  "request_id": 5,
+  "response_id": 2,
+  "FileBuffer": {
+   "pathspec": {
+    "path": "foo.idx",
+    "accessor": "ntfs"
+   },
+   "size": 196,
+   "data": "eyJmaWxlX29mZnNldCI6MCwib3JpZ2luYWxfb2Zmc2V0IjowLCJmaWxlX2xlbmd0aCI6NiwibGVuZ3RoIjo2fQp7ImZpbGVfb2Zmc2V0Ijo2LCJvcmlnaW5hbF9vZmZzZXQiOjYsImZpbGVfbGVuZ3RoIjowLCJsZW5ndGgiOjZ9CnsiZmlsZV9vZmZzZXQiOjYsIm9yaWdpbmFsX29mZnNldCI6MTIsImZpbGVfbGVuZ3RoIjo2LCJsZW5ndGgiOjZ9Cg==",
+   "eof": true
+  }
+ }
+]