diff --git a/client/taskenv/env.go b/client/taskenv/env.go index 377de9d93a7b..62f6c8673fc6 100644 --- a/client/taskenv/env.go +++ b/client/taskenv/env.go @@ -490,22 +490,19 @@ func (b *Builder) Build() *TaskEnv { } // Clean keys (see #2405) + prefixesToClean := [...]string{AddrPrefix, IpPrefix, PortPrefix, HostPortPrefix, MetaPrefix} cleanedEnv := make(map[string]string, len(envMap)) for k, v := range envMap { - cleanedK := helper.CleanEnvVar(k, '_') - cleanedEnv[cleanedK] = v - } - - var cleanedDeviceEnvs map[string]string - if deviceEnvs != nil { - cleanedDeviceEnvs = make(map[string]string, len(deviceEnvs)) - for k, v := range deviceEnvs { - cleanedK := helper.CleanEnvVar(k, '_') - cleanedDeviceEnvs[cleanedK] = v + cleanedK := k + for i := range prefixesToClean { + if strings.HasPrefix(k, prefixesToClean[i]) { + cleanedK = helper.CleanEnvVar(k, '_') + } } + cleanedEnv[cleanedK] = v } - return NewTaskEnv(cleanedEnv, cleanedDeviceEnvs, nodeAttrs) + return NewTaskEnv(cleanedEnv, deviceEnvs, nodeAttrs) } // Update task updates the environment based on a new alloc and task. diff --git a/client/taskenv/env_test.go b/client/taskenv/env_test.go index 0f2c280129c8..e4f10129966e 100644 --- a/client/taskenv/env_test.go +++ b/client/taskenv/env_test.go @@ -660,11 +660,17 @@ func TestEnvironment_AppendHostEnvvars(t *testing.T) { 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["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) } }