From dd864b1bace26f4387e0dd9a07bce86e83c7f1b1 Mon Sep 17 00:00:00 2001
From: Sebastian Rivera <sebastian.rivera@hashicorp.com>
Date: Tue, 5 Apr 2022 18:27:39 -0400
Subject: [PATCH] Add TF_WORKSPACE validation method

---
 internal/cloud/backend.go      | 76 ++++++++++++++++++++++++++++++++++
 internal/cloud/backend_test.go |  2 +-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/internal/cloud/backend.go b/internal/cloud/backend.go
index 08fd86eb30..ca895df46f 100644
--- a/internal/cloud/backend.go
+++ b/internal/cloud/backend.go
@@ -298,6 +298,16 @@ func (b *Cloud) Configure(obj cty.Value) tfdiags.Diagnostics {
 		return diags
 	}
 
+	if ws, ok := os.LookupEnv("TF_WORKSPACE"); ok {
+		if ws == b.WorkspaceMapping.Name || b.WorkspaceMapping.Strategy() == WorkspaceTagsStrategy {
+			diag := b.validWorkspaceEnvVar(context.Background(), b.organization, ws)
+			if diag != nil {
+				diags = diags.Append(diag)
+				return diags
+			}
+		}
+	}
+
 	// Check for the minimum version of Terraform Enterprise required.
 	//
 	// For API versions prior to 2.3, RemoteAPIVersion will return an empty string,
@@ -994,6 +1004,72 @@ func (b *Cloud) fetchWorkspace(ctx context.Context, organization string, workspa
 	return w, nil
 }
 
+// validWorkspaceEnvVar ensures we have selected a valid workspace using TF_WORKSPACE:
+// First, it ensures the workspace specified by TF_WORKSPACE exists in the organization
+// Second, if tags are specified in the configuration, it ensures TF_WORKSPACE belongs to the set
+// of available workspaces with those given tags.
+func (b *Cloud) validWorkspaceEnvVar(ctx context.Context, organization, workspace string) tfdiags.Diagnostic {
+	// first ensure the workspace exists
+	_, err := b.client.Workspaces.Read(ctx, organization, workspace)
+	if err != nil && err != tfe.ErrResourceNotFound {
+		return tfdiags.Sourceless(
+			tfdiags.Error,
+			"Terraform Cloud returned an unexpected error",
+			err.Error(),
+		)
+	}
+
+	if err == tfe.ErrResourceNotFound {
+		return tfdiags.Sourceless(
+			tfdiags.Error,
+			"Invalid workspace selection",
+			fmt.Sprintf(`Terraform failed to find workspace %q in organization %s.`, workspace, organization),
+		)
+	}
+
+	// if the configuration has specified tags, we need to ensure TF_WORKSPACE
+	// is a valid member
+	if b.WorkspaceMapping.Strategy() == WorkspaceTagsStrategy {
+		opts := &tfe.WorkspaceListOptions{}
+		opts.Tags = strings.Join(b.WorkspaceMapping.Tags, ",")
+
+		for {
+			wl, err := b.client.Workspaces.List(ctx, b.organization, opts)
+			if err != nil {
+				return tfdiags.Sourceless(
+					tfdiags.Error,
+					"Terraform Cloud returned an unexpected error",
+					err.Error(),
+				)
+			}
+
+			for _, ws := range wl.Items {
+				if ws.Name == workspace {
+					return nil
+				}
+			}
+
+			if wl.CurrentPage >= wl.TotalPages {
+				break
+			}
+
+			opts.PageNumber = wl.NextPage
+		}
+
+		return tfdiags.Sourceless(
+			tfdiags.Error,
+			"Invalid workspace selection",
+			fmt.Sprintf(
+				"Terraform failed to find workspace %q with the tags specified in your configuration:\n[%s]",
+				workspace,
+				strings.ReplaceAll(opts.Tags, ",", ", "),
+			),
+		)
+	}
+
+	return nil
+}
+
 func (wm WorkspaceMapping) tfeTags() []*tfe.Tag {
 	var tags []*tfe.Tag
 
diff --git a/internal/cloud/backend_test.go b/internal/cloud/backend_test.go
index f8799c7b62..9b0c0fbad8 100644
--- a/internal/cloud/backend_test.go
+++ b/internal/cloud/backend_test.go
@@ -303,7 +303,7 @@ func TestCloud_configWithEnvVars(t *testing.T) {
 			vars: map[string]string{
 				"TF_WORKSPACE": "i-dont-exist-in-org",
 			},
-			expectedErr: `Invalid workspace selection: "i-dont-exist-in-org" was not found in organization hashicorp`,
+			expectedErr: `Invalid workspace selection: Terraform failed to find workspace "i-dont-exist-in-org" in organization hashicorp`,
 		},
 		"workspaces and env var specified": {
 			config: cty.ObjectVal(map[string]cty.Value{