diff --git a/client/taskenv/env.go b/client/taskenv/env.go index 5b42a17a753..d0a1f7468f5 100644 --- a/client/taskenv/env.go +++ b/client/taskenv/env.go @@ -474,22 +474,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 531d8f16840..faadec1d3f0 100644 --- a/client/taskenv/env_test.go +++ b/client/taskenv/env_test.go @@ -656,11 +656,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) } }