From 2452f32476498116c5c395012287ecfbf96c90d0 Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Wed, 8 Mar 2017 16:44:46 -0800
Subject: [PATCH] Round two of env var cleaning

Should bring us into conformance with IEEE Std 1003.1, 2004 Edition:
http://pubs.opengroup.org/onlinepubs/009695399/basedefs/xbd_chap08.html

1 alloc/op and ~80ns/op on my machine.
---
 client/driver/env/env.go |  3 ++-
 helper/funcs.go          | 18 ++++++++++++++++++
 helper/funcs_test.go     | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/client/driver/env/env.go b/client/driver/env/env.go
index d496209bd..0887a0e8f 100644
--- a/client/driver/env/env.go
+++ b/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
diff --git a/helper/funcs.go b/helper/funcs.go
index 52ecd0f4f..8c62e06f2 100644
--- a/helper/funcs.go
+++ b/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)
+}
diff --git a/helper/funcs_test.go b/helper/funcs_test.go
index 200a8f490..507e948ae 100644
--- a/helper/funcs_test.go
+++ 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)
+	}
+}