diff --git a/client/taskenv/env.go b/client/taskenv/env.go index 5b42a17a7..30c88f1b8 100644 --- a/client/taskenv/env.go +++ b/client/taskenv/env.go @@ -476,7 +476,12 @@ func (b *Builder) Build() *TaskEnv { // Clean keys (see #2405) cleanedEnv := make(map[string]string, len(envMap)) for k, v := range envMap { - cleanedK := helper.CleanEnvVar(k, '_') + var cleanedK string + if strings.HasPrefix(k, "NOMAD_") { + cleanedK = helper.CleanEnvVar(k, '_', "") + } else { + cleanedK = helper.CleanEnvVar(k, '_', "-") + } cleanedEnv[cleanedK] = v } @@ -484,7 +489,12 @@ func (b *Builder) Build() *TaskEnv { if deviceEnvs != nil { cleanedDeviceEnvs = make(map[string]string, len(deviceEnvs)) for k, v := range deviceEnvs { - cleanedK := helper.CleanEnvVar(k, '_') + var cleanedK string + if strings.HasPrefix(k, "NOMAD_") { + cleanedK = helper.CleanEnvVar(k, '_', "") + } else { + cleanedK = helper.CleanEnvVar(k, '_', "-") + } cleanedDeviceEnvs[cleanedK] = v } } diff --git a/client/taskenv/env_test.go b/client/taskenv/env_test.go index fc102944b..faadec1d3 100644 --- a/client/taskenv/env_test.go +++ b/client/taskenv/env_test.go @@ -650,17 +650,24 @@ func TestEnvironment_AppendHostEnvvars(t *testing.T) { } } -// TestEnvironment_DashesInTaskName asserts dashes are kept in environment variables. -// See: https://github.com/hashicorp/nomad/issues/6079 +// TestEnvironment_DashesInTaskName asserts dashes in port labels are properly +// converted to underscores in environment variables. +// See: https://github.com/hashicorp/nomad/issues/2405 func TestEnvironment_DashesInTaskName(t *testing.T) { a := mock.Alloc() task := a.Job.TaskGroups[0].Tasks[0] - task.Env = map[string]string{"test-one-two": "three-four"} + task.Env = map[string]string{ + "test-one-two": "three-four", + "NOMAD_test_one_two": "three-five", + } envMap := NewBuilder(mock.Node(), a, task, "global").Build().Map() if envMap["test-one-two"] != "three-four" { t.Fatalf("Expected test-one-two=three-four in TaskEnv; found:\n%#v", envMap) } + if envMap["NOMAD_test_one_two"] != "three-five" { + t.Fatalf("Expected NOMAD_test_one_two=three-five in TaskEnv; found:\n%#v", envMap) + } } // TestEnvironment_UpdateTask asserts env vars and task meta are updated when a diff --git a/helper/funcs.go b/helper/funcs.go index 2a7a31b54..f17eb05c4 100644 --- a/helper/funcs.go +++ b/helper/funcs.go @@ -1,6 +1,7 @@ package helper import ( + "bytes" "crypto/sha512" "fmt" "regexp" @@ -343,11 +344,12 @@ func CopySliceInt(s []int) []int { // CleanEnvVar replaces all occurrences of illegal characters in an environment // variable with the specified byte. -func CleanEnvVar(s string, r byte) string { +// any extra characters will be considered valid +func CleanEnvVar(s string, r byte, e string) string { b := []byte(s) for i, c := range b { switch { - case c == '-': + case bytes.ContainsAny([]byte{c}, e): case c == '_': case c == '.': case c >= 'a' && c <= 'z': diff --git a/helper/funcs_test.go b/helper/funcs_test.go index a334ddde8..8a16cc1d1 100644 --- a/helper/funcs_test.go +++ b/helper/funcs_test.go @@ -135,13 +135,14 @@ func TestClearEnvVar(t *testing.T) { {"0sdf", "_sdf"}, {"asd0", "asd0"}, {"_asd", "_asd"}, - {"-asd", "-asd"}, + {"-asd", "_asd"}, + {"NOMAD_aaa-asd", "NOMAD_aaa_asd"}, {"asd.fgh", "asd.fgh"}, - {"A~!@#$%^&*()_+={}[]|\\;:'\"<,>?/Z", "A_____________________________Z"}, + {"A~!@#$%^&*()_+-={}[]|\\;:'\"<,>?/Z", "A______________________________Z"}, {"A\U0001f4a9Z", "A____Z"}, } for _, c := range cases { - if output := CleanEnvVar(c.input, '_'); output != c.expected { + if output := CleanEnvVar(c.input, '_', ""); output != c.expected { t.Errorf("CleanEnvVar(%q, '_') -> %q != %q", c.input, output, c.expected) } } @@ -154,6 +155,40 @@ func BenchmarkCleanEnvVar(b *testing.B) { b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { - CleanEnvVar(in, replacement) + CleanEnvVar(in, replacement, "") + } +} + +func TestClearEnvVarWithExtraChars(t *testing.T) { + type testCase struct { + input string + expected string + } + cases := []testCase{ + {"asdf", "asdf"}, + {"ASDF", "ASDF"}, + {"0sdf", "_sdf"}, + {"asd0", "asd0"}, + {"_asd", "_asd"}, + {"-asd", "-asd"}, + {"asd.fgh", "asd.fgh"}, + {"A~!@#$%^&*()_+-={}[]|\\;:'\"<,>?/Z", "A_____________-________________Z"}, + {"A\U0001f4a9Z", "A____Z"}, + } + for _, c := range cases { + if output := CleanEnvVar(c.input, '_', "-"); output != c.expected { + t.Errorf("CleanEnvVar(%q, '_') -> %q != %q", c.input, output, c.expected) + } + } +} + +func BenchmarkCleanEnvVarWithExtraChars(b *testing.B) { + in := "NOMAD_ADDR_redis-cache" + replacement := byte('_') + b.SetBytes(int64(len(in))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + CleanEnvVar(in, replacement, "-") } }