Merge pull request #1562 from snyk/fix/tf-hcl-schema

Support tfstate discovering within workspaces

pkg/cmd/scan.go

@@ -404,7 +404,8 @@ func retrieveBackendsFromHCL(workdir string) ([]config.SupplierConfig, error) {
 			continue
 		}
 
-		if supplierConfig := body.Backend.SupplierConfig(); supplierConfig != nil {
+		ws := hcl.GetCurrentWorkspaceName(path.Dir(match))
+		if supplierConfig := body.Backend.SupplierConfig(ws); supplierConfig != nil {
 			globaloutput.Printf(color.WhiteString("Using Terraform state %s found in %s. Use the --from flag to specify another state file.\n"), supplierConfig, match)
 			supplierConfigs = append(supplierConfigs, *supplierConfig)
 		}

pkg/terraform/hcl/backend.go

@@ -11,27 +11,28 @@ import (
 )
 
 type BackendBlock struct {
-	Name          string   `hcl:"name,label"`
-	Path          string   `hcl:"path,optional"`
-	WorkspaceDir  string   `hcl:"workspace_dir,optional"`
-	Bucket        string   `hcl:"bucket,optional"`
-	Key           string   `hcl:"key,optional"`
-	Region        string   `hcl:"region,optional"`
-	Prefix        string   `hcl:"prefix,optional"`
-	ContainerName string   `hcl:"container_name,optional"`
-	Remain        hcl.Body `hcl:",remain"`
+	Name               string   `hcl:"name,label"`
+	Path               string   `hcl:"path,optional"`
+	WorkspaceDir       string   `hcl:"workspace_dir,optional"`
+	Bucket             string   `hcl:"bucket,optional"`
+	Key                string   `hcl:"key,optional"`
+	Region             string   `hcl:"region,optional"`
+	Prefix             string   `hcl:"prefix,optional"`
+	ContainerName      string   `hcl:"container_name,optional"`
+	WorkspaceKeyPrefix string   `hcl:"workspace_key_prefix,optional"`
+	Remain             hcl.Body `hcl:",remain"`
 }
 
-func (b BackendBlock) SupplierConfig() *config.SupplierConfig {
+func (b BackendBlock) SupplierConfig(workspace string) *config.SupplierConfig {
 	switch b.Name {
 	case "local":
 		return b.parseLocalBackend()
 	case "s3":
-		return b.parseS3Backend()
+		return b.parseS3Backend(workspace)
 	case "gcs":
-		return b.parseGCSBackend()
+		return b.parseGCSBackend(workspace)
 	case "azurerm":
-		return b.parseAzurermBackend()
+		return b.parseAzurermBackend(workspace)
 	}
 	return nil
 }
@@ -47,32 +48,44 @@ func (b BackendBlock) parseLocalBackend() *config.SupplierConfig {
 	}
 }
 
-func (b BackendBlock) parseS3Backend() *config.SupplierConfig {
+func (b BackendBlock) parseS3Backend(ws string) *config.SupplierConfig {
 	if b.Bucket == "" || b.Key == "" {
 		return nil
 	}
+
+	keyPrefix := b.WorkspaceKeyPrefix
+	if ws != DefaultStateName {
+		if b.WorkspaceKeyPrefix == "" {
+			b.WorkspaceKeyPrefix = "env:"
+		}
+		keyPrefix = path.Join(b.WorkspaceKeyPrefix, ws)
+	}
+
 	return &config.SupplierConfig{
 		Key:     state.TerraformStateReaderSupplier,
 		Backend: backend.BackendKeyS3,
-		Path:    path.Join(b.Bucket, b.Key),
+		Path:    path.Join(b.Bucket, keyPrefix, b.Key),
 	}
 }
 
-func (b BackendBlock) parseGCSBackend() *config.SupplierConfig {
+func (b BackendBlock) parseGCSBackend(ws string) *config.SupplierConfig {
 	if b.Bucket == "" || b.Prefix == "" {
 		return nil
 	}
 	return &config.SupplierConfig{
 		Key:     state.TerraformStateReaderSupplier,
 		Backend: backend.BackendKeyGS,
-		Path:    fmt.Sprintf("%s.tfstate", path.Join(b.Bucket, b.Prefix)),
+		Path:    fmt.Sprintf("%s.tfstate", path.Join(b.Bucket, b.Prefix, ws)),
 	}
 }
 
-func (b BackendBlock) parseAzurermBackend() *config.SupplierConfig {
+func (b BackendBlock) parseAzurermBackend(ws string) *config.SupplierConfig {
 	if b.ContainerName == "" || b.Key == "" {
 		return nil
 	}
+	if ws != DefaultStateName {
+		b.Key = fmt.Sprintf("%senv:%s", b.Key, ws)
+	}
 	return &config.SupplierConfig{
 		Key:     state.TerraformStateReaderSupplier,
 		Backend: backend.BackendKeyAzureRM,

pkg/terraform/hcl/backend_test.go

@@ -1,78 +1,129 @@
 package hcl
 
 import (
+	"path"
 	"testing"
 
 	"github.com/snyk/driftctl/pkg/iac/config"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestHCL_getCurrentWorkspaceName(t *testing.T) {
+	cases := []struct {
+		name string
+		dir  string
+		want string
+	}{
+		{
+			name: "test with non-default workspace",
+			dir:  "testdata/foo_workspace",
+			want: "foo",
+		},
+		{
+			name: "test with non-existing directory",
+			dir:  "testdata/noenvfile",
+			want: "default",
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			workspace := GetCurrentWorkspaceName(tt.dir)
+			assert.Equal(t, tt.want, workspace)
+		})
+	}
+}
+
 func TestBackend_SupplierConfig(t *testing.T) {
 	cases := []struct {
-		name    string
-		dir     string
-		want    *config.SupplierConfig
-		wantErr string
+		name     string
+		filename string
+		want     *config.SupplierConfig
+		wantErr  string
 	}{
 		{
-			name:    "test with no backend block",
-			dir:     "testdata/no_backend_block.tf",
-			want:    nil,
-			wantErr: "testdata/no_backend_block.tf:1,11-11: Missing backend block; A backend block is required.",
+			name:     "test with no backend block",
+			filename: "testdata/no_backend_block.tf",
+			want:     nil,
+			wantErr:  "testdata/no_backend_block.tf:1,11-11: Missing backend block; A backend block is required.",
 		},
 		{
-			name: "test with local backend block",
-			dir:  "testdata/local_backend_block.tf",
+			name:     "test with local backend block",
+			filename: "testdata/local_backend_block.tf",
 			want: &config.SupplierConfig{
 				Key:  "tfstate",
 				Path: "terraform-state-prod/network/terraform.tfstate",
 			},
 		},
 		{
-			name: "test with S3 backend block",
-			dir:  "testdata/s3_backend_block.tf",
+			name:     "test with S3 backend block",
+			filename: "testdata/s3_backend_block.tf",
 			want: &config.SupplierConfig{
 				Key:     "tfstate",
 				Backend: "s3",
 				Path:    "terraform-state-prod/network/terraform.tfstate",
 			},
 		},
 		{
-			name: "test with GCS backend block",
-			dir:  "testdata/gcs_backend_block.tf",
+			name:     "test with S3 backend block with non-default workspace",
+			filename: "testdata/s3_backend_workspace/s3_backend_block.tf",
+			want: &config.SupplierConfig{
+				Key:     "tfstate",
+				Backend: "s3",
+				Path:    "terraform-state-prod/env:/bar/network/terraform.tfstate",
+			},
+		},
+		{
+			name:     "test with GCS backend block",
+			filename: "testdata/gcs_backend_block.tf",
 			want: &config.SupplierConfig{
 				Key:     "tfstate",
 				Backend: "gs",
-				Path:    "tf-state-prod/terraform/state.tfstate",
+				Path:    "tf-state-prod/terraform/state/default.tfstate",
 			},
 		},
 		{
-			name: "test with Azure backend block",
-			dir:  "testdata/azurerm_backend_block.tf",
+			name:     "test with Azure backend block",
+			filename: "testdata/azurerm_backend_block.tf",
 			want: &config.SupplierConfig{
 				Key:     "tfstate",
 				Backend: "azurerm",
 				Path:    "states/prod.terraform.tfstate",
 			},
 		},
+		{
+			name:     "test with Azure backend block with non-default workspace",
+			filename: "testdata/azurerm_backend_workspace/azurerm_backend_block.tf",
+			want: &config.SupplierConfig{
+				Key:     "tfstate",
+				Backend: "azurerm",
+				Path:    "states/prod.terraform.tfstateenv:bar",
+			},
+		},
+		{
+			name:     "test with unknown backend",
+			filename: "testdata/unknown_backend_block.tf",
+			want:     nil,
+		},
 	}
 
 	for _, tt := range cases {
 		t.Run(tt.name, func(t *testing.T) {
-			hcl, err := ParseTerraformFromHCL(tt.dir)
+			hcl, err := ParseTerraformFromHCL(tt.filename)
 			if tt.wantErr == "" {
 				assert.NoError(t, err)
 			} else {
 				assert.EqualError(t, err, tt.wantErr)
 				return
 			}
 
-			if hcl.Backend.SupplierConfig() == nil {
+			ws := GetCurrentWorkspaceName(path.Dir(tt.filename))
+			if hcl.Backend.SupplierConfig(ws) == nil {
 				assert.Nil(t, tt.want)
 				return
 			}
 
-			assert.Equal(t, *tt.want, *hcl.Backend.SupplierConfig())
+			assert.Equal(t, *tt.want, *hcl.Backend.SupplierConfig(ws))
 		})
 	}
 }

pkg/terraform/hcl/hcl.go

@@ -1,31 +1,53 @@
 package hcl
 
 import (
+	"io/ioutil"
+	"path"
+	"strings"
+
+	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/gohcl"
 	"github.com/hashicorp/hcl/v2/hclparse"
 )
 
+const DefaultStateName = "default"
+
 type MainBodyBlock struct {
 	Terraform TerraformBlock `hcl:"terraform,block"`
+	Remain    hcl.Body       `hcl:",remain"`
 }
 
 type TerraformBlock struct {
 	Backend BackendBlock `hcl:"backend,block"`
+	Remain  hcl.Body     `hcl:",remain"`
 }
 
 func ParseTerraformFromHCL(filename string) (*TerraformBlock, error) {
-	var v MainBodyBlock
+	var body MainBodyBlock
 
 	parser := hclparse.NewParser()
 	f, diags := parser.ParseHCLFile(filename)
 	if diags.HasErrors() {
 		return nil, diags
 	}
 
-	diags = gohcl.DecodeBody(f.Body, nil, &v)
+	diags = gohcl.DecodeBody(f.Body, nil, &body)
 	if diags.HasErrors() {
 		return nil, diags
 	}
 
-	return &v.Terraform, nil
+	return &body.Terraform, nil
+}
+
+func GetCurrentWorkspaceName(cwd string) string {
+	name := DefaultStateName // See https://github.com/hashicorp/terraform/blob/main/internal/backend/backend.go#L33
+
+	data, err := ioutil.ReadFile(path.Join(cwd, ".terraform/environment"))
+	if err != nil {
+		return name
+	}
+	if v := strings.Trim(string(data), "\n"); v != "" {
+		name = v
+	}
+	return name
 }

pkg/terraform/hcl/testdata/azurerm_backend_block.tf

@@ -6,3 +6,7 @@ terraform {
         key                  = "prod.terraform.tfstate"
     }
 }
+
+provider "azurerm" {
+    features {}
+}

pkg/terraform/hcl/testdata/azurerm_backend_workspace/.gitignore

@@ -0,0 +1 @@
+!.terraform

pkg/terraform/hcl/testdata/azurerm_backend_workspace/.terraform/environment

@@ -0,0 +1 @@
+bar

pkg/terraform/hcl/testdata/azurerm_backend_workspace/azurerm_backend_block.tf

@@ -0,0 +1,12 @@
+terraform {
+    backend "azurerm" {
+        resource_group_name  = "StorageAccount-ResourceGroup"
+        storage_account_name = "abcd1234"
+        container_name       = "states"
+        key                  = "prod.terraform.tfstate"
+    }
+}
+
+provider "azurerm" {
+    features {}
+}

pkg/terraform/hcl/testdata/foo_workspace/.gitignore

@@ -0,0 +1 @@
+!.terraform

pkg/terraform/hcl/testdata/foo_workspace/.terraform/environment

@@ -0,0 +1 @@
+foo

pkg/terraform/hcl/testdata/gcs_backend_block.tf

@@ -4,3 +4,9 @@ terraform {
         prefix  = "terraform/state"
     }
 }
+
+provider "google" {
+    project = "my-project"
+    region  = "us-central1"
+    zone    = "us-central1-c"
+}

pkg/terraform/hcl/testdata/s3_backend_block.tf

@@ -5,3 +5,5 @@ terraform {
         region = "us-east-1"
     }
 }
+
+provider "aws" {}

pkg/terraform/hcl/testdata/s3_backend_workspace/.gitignore

@@ -0,0 +1 @@
+!.terraform

pkg/terraform/hcl/testdata/s3_backend_workspace/.terraform/environment

@@ -0,0 +1 @@
+bar

pkg/terraform/hcl/testdata/s3_backend_workspace/s3_backend_block.tf

@@ -0,0 +1,9 @@
+terraform {
+    backend "s3" {
+        bucket = "terraform-state-prod"
+        key    = "network/terraform.tfstate"
+        region = "us-east-1"
+    }
+}
+
+provider "aws" {}

pkg/terraform/hcl/testdata/unknown_backend_block.tf

@@ -0,0 +1,3 @@
+terraform {
+    backend "oss" {}
+}