mirror of
https://github.com/empayre/fleet.git
synced 2024-11-07 01:15:22 +00:00
c89cd370d5
This adds the option to set up an S3 bucket as the storage backend for file carving (partially solving #111). It works by using the multipart upload capabilities of S3 to maintain compatibility with the "upload in blocks" protocol that osquery uses. It does this basically replacing the carve_blocks table while still maintaining the metadata in the original place (it would probably be possible to rely completely on S3 by using object tagging at the cost of listing performance). To make this pluggable, I created a new field in the service struct dedicated to the CarveStore which, if no configuration for S3 is set up will be just a reference to the standard datastore, otherwise it will point to the S3 one (effectively this separation will allow in the future to add more backends).
442 lines
12 KiB
Go
442 lines
12 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
hostctx "github.com/fleetdm/fleet/server/contexts/host"
|
|
"github.com/fleetdm/fleet/server/kolide"
|
|
"github.com/fleetdm/fleet/server/mock"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestCarveBegin(t *testing.T) {
|
|
host := kolide.Host{ID: 3}
|
|
payload := kolide.CarveBeginPayload{
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
expectedMetadata := kolide.CarveMetadata{
|
|
ID: 7,
|
|
HostId: host.ID,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
}
|
|
ms.NewCarveFunc = func(metadata *kolide.CarveMetadata) (*kolide.CarveMetadata, error) {
|
|
metadata.ID = 7
|
|
return metadata, nil
|
|
}
|
|
|
|
ctx := hostctx.NewContext(context.Background(), host)
|
|
|
|
metadata, err := svc.CarveBegin(ctx, payload)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, metadata.SessionId)
|
|
metadata.SessionId = "" // Clear this before comparison
|
|
metadata.Name = "" // Clear this before comparison
|
|
metadata.CreatedAt = time.Time{} // Clear this before comparison
|
|
assert.Equal(t, expectedMetadata, *metadata)
|
|
}
|
|
|
|
func TestCarveBeginNewCarveError(t *testing.T) {
|
|
host := kolide.Host{ID: 3}
|
|
payload := kolide.CarveBeginPayload{
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.NewCarveFunc = func(metadata *kolide.CarveMetadata) (*kolide.CarveMetadata, error) {
|
|
return nil, fmt.Errorf("ouch!")
|
|
}
|
|
|
|
ctx := hostctx.NewContext(context.Background(), host)
|
|
|
|
_, err := svc.CarveBegin(ctx, payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ouch!")
|
|
}
|
|
|
|
func TestCarveBeginEmptyError(t *testing.T) {
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ctx := hostctx.NewContext(context.Background(), kolide.Host{})
|
|
|
|
_, err := svc.CarveBegin(ctx, kolide.CarveBeginPayload{})
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "carve_size must be greater than 0")
|
|
}
|
|
|
|
func TestCarveBeginMissingHostError(t *testing.T) {
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
|
|
_, err := svc.CarveBegin(context.Background(), kolide.CarveBeginPayload{})
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "missing host")
|
|
}
|
|
|
|
func TestCarveBeginBlockSizeMaxError(t *testing.T) {
|
|
host := kolide.Host{ID: 3}
|
|
payload := kolide.CarveBeginPayload{
|
|
BlockCount: 10,
|
|
BlockSize: 1024 * 1024 * 1024 * 1024, // 1TB
|
|
CarveSize: 10 * 1024 * 1024 * 1024 * 1024, // 10TB
|
|
RequestId: "carve_request",
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
|
|
ctx := hostctx.NewContext(context.Background(), host)
|
|
|
|
_, err := svc.CarveBegin(ctx, payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "block_size exceeds max")
|
|
}
|
|
|
|
func TestCarveBeginCarveSizeMaxError(t *testing.T) {
|
|
host := kolide.Host{ID: 3}
|
|
payload := kolide.CarveBeginPayload{
|
|
BlockCount: 1024 * 1024,
|
|
BlockSize: 10 * 1024 * 1024, // 1TB
|
|
CarveSize: 10 * 1024 * 1024 * 1024 * 1024, // 10TB
|
|
RequestId: "carve_request",
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
|
|
ctx := hostctx.NewContext(context.Background(), host)
|
|
|
|
_, err := svc.CarveBegin(ctx, payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "carve_size exceeds max")
|
|
}
|
|
|
|
func TestCarveBeginCarveSizeError(t *testing.T) {
|
|
host := kolide.Host{ID: 3}
|
|
payload := kolide.CarveBeginPayload{
|
|
BlockCount: 7,
|
|
BlockSize: 13,
|
|
CarveSize: 7*13 + 1,
|
|
RequestId: "carve_request",
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ctx := hostctx.NewContext(context.Background(), host)
|
|
|
|
// Too big
|
|
_, err := svc.CarveBegin(ctx, payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "carve_size does not match")
|
|
|
|
// Too small
|
|
payload.CarveSize = 6 * 13
|
|
_, err = svc.CarveBegin(ctx, payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "carve_size does not match")
|
|
}
|
|
|
|
func TestCarveCarveBlockGetCarveError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
return nil, fmt.Errorf("ouch!")
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
SessionId: sessionId,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ouch!")
|
|
}
|
|
|
|
func TestCarveCarveBlockRequestIdError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
RequestId: "not_matching",
|
|
SessionId: sessionId,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "request_id does not match")
|
|
}
|
|
|
|
func TestCarveCarveBlockBlockCountExceedError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
BlockId: 23,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "block_id exceeds expected max")
|
|
}
|
|
|
|
func TestCarveCarveBlockBlockCountMatchError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
BlockId: 7,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "block_id does not match")
|
|
}
|
|
|
|
func TestCarveCarveBlockBlockSizeError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 16,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :) TOO LONG!!!"),
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
BlockId: 4,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "exceeded declared block size")
|
|
}
|
|
|
|
func TestCarveCarveBlockNewBlockError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
ms.NewBlockFunc = func(carve *kolide.CarveMetadata, blockId int64, data []byte) error {
|
|
return fmt.Errorf("kaboom!")
|
|
}
|
|
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
BlockId: 4,
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "kaboom!")
|
|
}
|
|
|
|
func TestCarveCarveBlock(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
payload := kolide.CarveBlockPayload{
|
|
Data: []byte("this is the carve data :)"),
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
BlockId: 4,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveBySessionIdFunc = func(sessionId string) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.SessionId, sessionId)
|
|
return metadata, nil
|
|
}
|
|
ms.NewBlockFunc = func(carve *kolide.CarveMetadata, blockId int64, data []byte) error {
|
|
assert.Equal(t, metadata, carve)
|
|
assert.Equal(t, int64(4), blockId)
|
|
assert.Equal(t, payload.Data, data)
|
|
return nil
|
|
}
|
|
|
|
err := svc.CarveBlock(context.Background(), payload)
|
|
require.NoError(t, err)
|
|
assert.True(t, ms.NewBlockFuncInvoked)
|
|
}
|
|
|
|
func TestCarveGetBlock(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveFunc = func(carveId int64) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.ID, carveId)
|
|
return metadata, nil
|
|
}
|
|
ms.GetBlockFunc = func(carve *kolide.CarveMetadata, blockId int64) ([]byte, error) {
|
|
assert.Equal(t, metadata.ID, carve.ID)
|
|
assert.Equal(t, int64(3), blockId)
|
|
return []byte("foobar"), nil
|
|
}
|
|
|
|
data, err := svc.GetBlock(context.Background(), metadata.ID, 3)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, []byte("foobar"), data)
|
|
}
|
|
|
|
func TestCarveGetBlockNotAvailableError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveFunc = func(carveId int64) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.ID, carveId)
|
|
return metadata, nil
|
|
}
|
|
|
|
// Block requested is great than max block
|
|
_, err := svc.GetBlock(context.Background(), metadata.ID, 7)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "not yet available")
|
|
}
|
|
|
|
func TestCarveGetBlockGetBlockError(t *testing.T) {
|
|
sessionId := "foobar"
|
|
metadata := &kolide.CarveMetadata{
|
|
ID: 2,
|
|
HostId: 3,
|
|
BlockCount: 23,
|
|
BlockSize: 64,
|
|
CarveSize: 23 * 64,
|
|
RequestId: "carve_request",
|
|
SessionId: sessionId,
|
|
MaxBlock: 3,
|
|
}
|
|
ms := new(mock.Store)
|
|
svc := service{carveStore: ms}
|
|
ms.CarveFunc = func(carveId int64) (*kolide.CarveMetadata, error) {
|
|
assert.Equal(t, metadata.ID, carveId)
|
|
return metadata, nil
|
|
}
|
|
ms.GetBlockFunc = func(carve *kolide.CarveMetadata, blockId int64) ([]byte, error) {
|
|
assert.Equal(t, metadata.ID, carve.ID)
|
|
assert.Equal(t, int64(3), blockId)
|
|
return nil, fmt.Errorf("yow!!")
|
|
}
|
|
|
|
// Block requested is great than max block
|
|
_, err := svc.GetBlock(context.Background(), metadata.ID, 3)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "yow!!")
|
|
}
|