From f64ade2304583ba338c4da3c6c11de722ea4b497 Mon Sep 17 00:00:00 2001
From: James Rasell <jrasell@users.noreply.github.com>
Date: Tue, 24 Oct 2023 16:48:13 +0100
Subject: [PATCH] cli: ensure HCL env vars are added to the job submission
 object. (#18832)

---
 .changelog/18832.txt    |  3 +++
 command/helpers.go      | 50 ++++++++++++++++++++++++++++++++++++++---
 command/helpers_test.go | 48 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 3 deletions(-)
 create mode 100644 .changelog/18832.txt

diff --git a/.changelog/18832.txt b/.changelog/18832.txt
new file mode 100644
index 000000000..2ab3a493b
--- /dev/null
+++ b/.changelog/18832.txt
@@ -0,0 +1,3 @@
+```release-note:bug
+cli: ensure HCL env vars are added to the job submission object in the `job run` command
+```
diff --git a/command/helpers.go b/command/helpers.go
index f7b7b54cb..7453978fa 100644
--- a/command/helpers.go
+++ b/command/helpers.go
@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"maps"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -534,6 +535,10 @@ func (j *JobGetter) Get(jpath string) (*api.JobSubmission, *api.Job, error) {
 			return nil, nil, fmt.Errorf("Failed to parse HCL job: %w", err)
 		}
 
+		// Perform the environment listing here as it is used twice beyond this
+		// point.
+		osEnv := os.Environ()
+
 		// we are parsing HCL2, whether from a file or stdio
 		jobStruct, err = jobspec2.ParseWithConfig(&jobspec2.ParseConfig{
 			Path:     pathName,
@@ -541,7 +546,7 @@ func (j *JobGetter) Get(jpath string) (*api.JobSubmission, *api.Job, error) {
 			ArgVars:  j.Vars,
 			AllowFS:  true,
 			VarFiles: j.VarFiles,
-			Envs:     os.Environ(),
+			Envs:     osEnv,
 			Strict:   j.Strict,
 		})
 
@@ -549,15 +554,24 @@ func (j *JobGetter) Get(jpath string) (*api.JobSubmission, *api.Job, error) {
 		var readVarFileErr error
 		if err == nil {
 			// combine any -var-file data into one big blob
-			varFileCat, readVarFileErr = extractVarFiles([]string(j.VarFiles))
+			varFileCat, readVarFileErr = extractVarFiles(j.VarFiles)
 			if readVarFileErr != nil {
 				return nil, nil, fmt.Errorf("Failed to read var file(s): %w", readVarFileErr)
 			}
 		}
 
+		// Extract variables declared by the -var flag and as environment
+		// variables.
+		extractedVarFlags := extractVarFlags(j.Vars)
+		extractedEnvVars := extractJobSpecEnvVars(osEnv)
+
+		// Merge the two maps ensuring that variables defined by -var flags
+		// take precedence.
+		maps.Copy(extractedEnvVars, extractedVarFlags)
+
 		// submit the job with the submission with content from -var flags
 		jobSubmission = &api.JobSubmission{
-			VariableFlags: extractVarFlags(j.Vars),
+			VariableFlags: extractedEnvVars,
 			Variables:     varFileCat,
 			Source:        source.String(),
 			Format:        formatHCL2,
@@ -606,6 +620,36 @@ func extractVarFlags(slice []string) map[string]string {
 	return m
 }
 
+// extractJobSpecEnvVars is used to extract Nomad specific HCL variables from
+// the OS environment. The input envVars parameter is expected to be generated
+// from the os.Environment function call. The result is never nil for
+// convenience.
+func extractJobSpecEnvVars(envVars []string) map[string]string {
+
+	m := make(map[string]string)
+
+	for _, raw := range envVars {
+		if !strings.HasPrefix(raw, jobspec2.VarEnvPrefix) {
+			continue
+		}
+
+		// Trim the prefix, so we just have the raw key=value variable
+		// remaining.
+		raw = raw[len(jobspec2.VarEnvPrefix):]
+
+		// Identify the index of the equals sign which is where we split the
+		// variable k/v pair. -1 indicates the equals sign is not found and
+		// therefore the var is not valid.
+		if eq := strings.Index(raw, "="); eq == -1 {
+			continue
+		} else if raw[:eq] != "" {
+			m[raw[:eq]] = raw[eq+1:]
+		}
+	}
+
+	return m
+}
+
 // mergeAutocompleteFlags is used to join multiple flag completion sets.
 func mergeAutocompleteFlags(flags ...complete.Flags) complete.Flags {
 	merged := make(map[string]complete.Predictor, len(flags))
diff --git a/command/helpers_test.go b/command/helpers_test.go
index 26cba2ada..25fb5b3b2 100644
--- a/command/helpers_test.go
+++ b/command/helpers_test.go
@@ -685,3 +685,51 @@ func Test_extractVarFlags(t *testing.T) {
 		}, result)
 	})
 }
+
+func Test_extractJobSpecEnvVars(t *testing.T) {
+	ci.Parallel(t)
+
+	t.Run("nil", func(t *testing.T) {
+		must.MapEmpty(t, extractJobSpecEnvVars(nil))
+	})
+
+	t.Run("complete", func(t *testing.T) {
+		result := extractJobSpecEnvVars([]string{
+			"NOMAD_VAR_count=13",
+			"GOPATH=/Users/jrasell/go",
+			"NOMAD_VAR_image=redis:7",
+		})
+		must.Eq(t, map[string]string{
+			"count": "13",
+			"image": "redis:7",
+		}, result)
+	})
+
+	t.Run("whitespace", func(t *testing.T) {
+		result := extractJobSpecEnvVars([]string{
+			"NOMAD_VAR_count = 13",
+			"GOPATH = /Users/jrasell/go",
+		})
+		must.Eq(t, map[string]string{
+			"count ": " 13",
+		}, result)
+	})
+
+	t.Run("empty key", func(t *testing.T) {
+		result := extractJobSpecEnvVars([]string{
+			"NOMAD_VAR_=13",
+			"=/Users/jrasell/go",
+		})
+		must.Eq(t, map[string]string{}, result)
+	})
+
+	t.Run("empty value", func(t *testing.T) {
+		result := extractJobSpecEnvVars([]string{
+			"NOMAD_VAR_count=",
+			"GOPATH=",
+		})
+		must.Eq(t, map[string]string{
+			"count": "",
+		}, result)
+	})
+}