Check if ECS data directory exists within agent container

Author: mye956

agent/data/client.go

@@ -14,6 +14,8 @@
 package data
 
 import (
+	"fmt"
+	"os"
 	"path/filepath"
 	"sync"
 
@@ -39,6 +41,7 @@ const (
 	resAttachmentsBucketName = "resattachments"
 	metadataBucketName       = "metadata"
 	emptyAgentVersionMsg     = "No version info available in boltDB. Either this is a fresh instance, or we were using state file to persist data. Transformer not applicable."
+	notDirectoryErrorMsg     = "path %s is not a valid directory"
 )
 
 var (
@@ -53,6 +56,7 @@ var (
 		resAttachmentsBucketName,
 		metadataBucketName,
 	}
+	dirExists = checkDirectoryExists
 )
 
 // Client specifies the data management interface to persist and manage various kinds of data in the agent.
@@ -133,7 +137,15 @@ func NewWithSetup(dataDir string) (Client, error) {
 // setup initiates the boltdb client and makes sure the buckets we use and transformer are created, and
 // registers transformation functions to transformer.
 func setup(dataDir string) (*client, error) {
+	// Check if the directory is valid
+	err := dirExists(dataDir)
+	if err != nil {
+		return nil, err
+	}
 	db, err := bolt.Open(filepath.Join(dataDir, dbName), dbMode, nil)
+	if err != nil {
+		return nil, err
+	}
 	err = db.Update(func(tx *bolt.Tx) error {
 		for _, b := range buckets {
 			_, err = tx.CreateBucketIfNotExists([]byte(b))
@@ -167,3 +179,15 @@ func setup(dataDir string) (*client, error) {
 func (c *client) Close() error {
 	return c.DB.Close()
 }
+
+// Checks if the data directory exists within the agent container
+func checkDirectoryExists(path string) error {
+	fileInfo, err := os.Stat(path)
+	if err != nil {
+		return err
+	}
+	if !fileInfo.IsDir() {
+		return fmt.Errorf(notDirectoryErrorMsg, path)
+	}
+	return nil
+}

agent/data/client_test.go

@@ -17,16 +17,22 @@
 package data
 
 import (
+	"fmt"
 	"path/filepath"
 	"testing"
 
 	generaldata "github.com/aws/amazon-ecs-agent/ecs-agent/data"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/modeltransformer"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	bolt "go.etcd.io/bbolt"
 )
 
+const (
+	statPathNotFoundErr = "stat %s: no such file or directory"
+)
+
 func newTestClient(t *testing.T) Client {
 	testDir := t.TempDir()
 
@@ -56,3 +62,43 @@ func newTestClient(t *testing.T) Client {
 	})
 	return testClient
 }
+
+func TestSetupInvalidDirectoryErrors(t *testing.T) {
+	originalDirExists := dirExists
+	defer func() {
+		dirExists = originalDirExists
+	}()
+	tcs := []struct {
+		name          string
+		dirPath       string
+		mockDirExists func(string) error
+		expectedError error
+	}{
+		{
+			name:    "path not found",
+			dirPath: "/path/",
+			mockDirExists: func(path string) error {
+				return fmt.Errorf(statPathNotFoundErr, path)
+			},
+			expectedError: fmt.Errorf(statPathNotFoundErr, "/path/"),
+		},
+		{
+			name:    "path is not a directory",
+			dirPath: "/data/blah.txt",
+			mockDirExists: func(path string) error {
+				return fmt.Errorf(notDirectoryErrorMsg, path)
+			},
+			expectedError: fmt.Errorf(notDirectoryErrorMsg, "/data/blah.txt"),
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			dirExists = tc.mockDirExists
+
+			_, err := setup(tc.dirPath)
+
+			assert.Equal(t, tc.expectedError, err)
+		})
+	}
+}