Merge pull request #2416 from hashicorp/b-clean-env-var-2

Round two of env var cleaning

client/driver/env/env.go

@@ -7,6 +7,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/hashicorp/nomad/helper"
 	hargs "github.com/hashicorp/nomad/helper/args"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
@@ -254,7 +255,7 @@ func (t *TaskEnvironment) Build() *TaskEnvironment {
 	// Clean keys (see #2405)
 	cleanedEnv := make(map[string]string, len(t.TaskEnv))
 	for k, v := range t.TaskEnv {
-		cleanedK := strings.Replace(k, "-", "_", -1)
+		cleanedK := helper.CleanEnvVar(k, '_')
 		cleanedEnv[cleanedK] = v
 	}
 	t.TaskEnv = cleanedEnv

helper/funcs.go

@@ -177,3 +177,21 @@ func CopySliceInt(s []int) []int {
 	}
 	return c
 }
+
+// CleanEnvVar replaces all occurrences of illegal characters in an environment
+// variable with the specified byte.
+func CleanEnvVar(s string, r byte) string {
+	b := []byte(s)
+	for i, c := range b {
+		switch {
+		case c == '_':
+		case c >= 'a' && c <= 'z':
+		case c >= 'A' && c <= 'Z':
+		case i > 0 && c >= '0' && c <= '9':
+		default:
+			// Replace!
+			b[i] = r
+		}
+	}
+	return string(b)
+}

helper/funcs_test.go

@@ -35,3 +35,36 @@ func TestMapStringStringSliceValueSet(t *testing.T) {
 		t.Fatalf("Bad; got %v; want %v", act, exp)
 	}
 }
+
+func TestClearEnvVar(t *testing.T) {
+	type testCase struct {
+		input    string
+		expected string
+	}
+	cases := []testCase{
+		{"asdf", "asdf"},
+		{"ASDF", "ASDF"},
+		{"0sdf", "_sdf"},
+		{"asd0", "asd0"},
+		{"_asd", "_asd"},
+		{"-asd", "_asd"},
+		{"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 BenchmarkCleanEnvVar(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)
+	}
+}

website/source/docs/runtime/environment.html.md

@@ -90,9 +90,9 @@ environment variables.
   </tr>
 </table>
 
-~> Port labels and task names will have any dashes `-` in their names replaced by
-underscores `_` when they're used in environment variable names such as
-`NOMAD_ADDR_<task>_<label>`.
+~> Port labels and task names will have any non-alphanumeric or underscore
+characters in their names replaced by underscores `_` when they're used in
+environment variable names such as `NOMAD_ADDR_<task>_<label>`.
 
 ## Task Identifiers