environ: Fix key case, add GetEmpty and Remove.

common/system/environ/environment.go

 // Copyright 2016 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 // Package environ is a simple environment variable manipulation library. It
 // couples the system's environment, representated as a []string of KEY[=VALUE]
 // strings, into a key/value map and back.
 package environ
 
 import (
         "os"
+        "runtime"
         "sort"
         "strings"
 )
 
+// SystemIsCaseInsensitive returns true if the local operating system treats
+// environment variable keys in a case-insensitive manner.
+func SystemIsCaseInsensitive() bool {
+        return runtime.GOOS == "windows"
+}
+
 // Env contains system environment variables. It preserves each environment
 // variable verbatim (even if invalid).
 type Env struct {
+        // CaseInsensitive is true if environment keys are case-insensitive.
+        CaseInsensitive bool
+
         // env is a map of envoironment key to its "KEY=VALUE" value.
         //
         // Note that the value here is the full "KEY=VALUE", not just the VALUE part.
         // This allows us to reconstitute the original environment string slice
         // without reallocating all of its composite strings.
         env map[string]string
 }
 
+// normalizeKey normalizes the map key, ensuring that it is handled
+// case-insensitive.
+func (e *Env) normalizeKey(k string) string {
+        if e.CaseInsensitive {
+                return strings.ToUpper(k)
+        }
+        return k
+}
+
 // System returns an Env instance instantiated with the current os.Environ
 // values.
+//
+// The environment is automatically configured with the local system's case
+// insensitivity.
 func System() Env {
         return New(os.Environ())
 }
 
 // New instantiates a new Env instance from the supplied set of environment
-// KEY=VALUE strings.
+// KEY=VALUE strings using LoadSlice.
+//
+// The environment is automatically configured with the local system's case
+// insensitivity.
 func New(s []string) Env {
-        e := Env{}
-        if len(s) > 0 {
-                e.env = make(map[string]string, len(s))
-                for _, v := range s {
-                        k, _ := Split(v)
-                        e.env[k] = v
-                }
+        e := Env{
+                CaseInsensitive: SystemIsCaseInsensitive(),
         }
+        e.LoadSlice(s)
         return e
 }
 
-// Make creaets a new Env from an environment key/value map.
-func Make(v map[string]string) Env {
-        e := Env{}
-        if len(v) > 0 {
-                e.env = make(map[string]string, len(v))
-                for k, v := range v {
-                        e.Set(k, v)
-                }
+// Load adds environment variables defined in a key/value map to an existing
+// environment.
+func (e *Env) Load(m map[string]string) {
+        if len(m) == 0 {
+                return
         }
-        return e
+        if e.env == nil {
+                e.env = make(map[string]string, len(m))
+        }
+        for k, v := range m {
+                e.Set(k, v)
+        }
 }
 
+// LoadSlice adds environment variable entries, specified as "KEY[=VALUE]"
+// strings, into an existing environment.
+//
+// Entries are added in order. If there are duplicates, later entries will
+// override earlier ones.
+func (e *Env) LoadSlice(s []string) {
+        if len(s) == 0 {
+                return
+        }
+        if e.env == nil {
+                e.env = make(map[string]string, len(s))
+        }
+        for _, entry := range s {
+                e.SetEntry(entry)
+        }
+}
+
+// Set sets the supplied environment key and value.
+func (e *Env) Set(k, v string) {
+        if e.env == nil {
+                e.env = make(map[string]string)
+        }
+        e.env[e.normalizeKey(k)] = Join(k, v)
+}
+
+// SetEntry sets the supplied environment to a "KEY[=VALUE]" entry.
+func (e *Env) SetEntry(entry string) { e.Set(Split(entry)) }
+
+// Remove removes a value from the environment, returning true if the value was
+// present. If the value was missing, Remove returns false.
+//
+// Remove is different from Set(k, "") in that Set persists the key with an
+// empty value, while Remove removes it entirely.
+func (e *Env) Remove(k string) bool {
+        k = e.normalizeKey(k)
+        if _, ok := e.env[k]; ok {
+                delete(e.env, k)
+                return true
+        }
+        return false
+}
+
+//
+// NOTE to implementers: all mutation methods MUST accept a pointer Env, as they
+// may mutate the underlying "env" map value.
+//
+// Read-only accessors should accept Env as a value.
+//
+
 // Get returns the environment value for the supplied key.
 //
 // If the value is defined, ok will be true and v will be set to its value (note
 // that this can be empty if the environment has an empty value). If the value
 // is not defined, ok will be false.
 func (e Env) Get(k string) (v string, ok bool) {
-        if v, ok = e.env[k]; ok {
-                v = value(k, v)
+        // NOTE: "v" is initially the combined "key=value" entry, and will need to be
+        // split.
+        if v, ok = e.env[e.normalizeKey(k)]; ok {
+                _, v = splitEntryGivenKey(k, v)
         }
         return
 }
 
-// Set sets the supplied environment key and value.
-func (e *Env) Set(k, v string) {
-        if e.env == nil {
-                e.env = make(map[string]string)
-        }
-        e.env[k] = Join(k, v)
+// GetEmpty is the same as Get, except that instead of returning a separate
+// boolean indicating the presence of a key, it will return an empty string if
+// the key is missing.
+func (e Env) GetEmpty(k string) string {
+        v, _ := e.Get(k)
+        return v
 }
 
 // Sorted returns the contents of the environment, sorted by key.
 func (e Env) Sorted() []string {
         var r []string
         if len(e.env) > 0 {
                 r = make([]string, 0, len(e.env))
                 for _, v := range e.env {
                         r = append(r, v)
                 }
                 sort.Strings(r)
         }
         return r
 }
 
 // Map returns a map of the key/value values in the environment.
 //
 // This is a clone of the contents of e; manipulating this map will not change
 // the values in e.
 func (e Env) Map() map[string]string {
         if len(e.env) == 0 {
                 return nil
         }
 
         m := make(map[string]string, len(e.env))
-        for k, v := range e.env {
-                m[k] = value(k, v)
+        for k, entry := range e.env {
+                ek, ev := splitEntryGivenKey(k, entry)
+                m[ek] = ev
         }
         return m
 }
 
 // Len returns the number of environment variables defined in e.
 func (e Env) Len() int { return len(e.env) }
 
 // Clone creates a new Env instance that is identical to, but independent from,
 // e.
 func (e Env) Clone() Env {
-        var clone Env
+        clone := e
+        clone.env = nil
+
         if len(e.env) > 0 {
                 clone.env = make(map[string]string, len(e.env))
                 for k, v := range e.env {
                         clone.env[k] = v
                 }
         }
         return clone
 }
 
 // Split splits the supplied environment variable value into a key/value pair.
 //
 // If v is of the form:
 //      - KEY, returns (KEY, "")
 //      - KEY=, returns (KEY, "")
 //      - KEY=VALUE, returns (KEY, VALUE)
 func Split(v string) (key, value string) {
         parts := strings.SplitN(v, "=", 2)
         switch len(parts) {
         case 2:
                 value = parts[1]
                 fallthrough
 
         case 1:
                 key = parts[0]
         }
         return
 }
 
-// getValue returns the value from an internal "env" map value field.
-//
-// It assumes that the environment variable is well-formed, one of:
-// - KEY
-// - KEY=
-// - KEY=VALUE
-func value(key, v string) string {
-        prefixSize := len(key) + len("=")
-        if len(v) <= prefixSize {
-                return ""
-        }
-        return v[prefixSize:]
-}
-
 // Join creates an environment variable definition for the supplied key/value.
 func Join(k, v string) string {
         if v == "" {
                 return k
         }
-        return strings.Join([]string{k, v}, "=")
+        return k + "=" + v
 }
+
+// splitEntryGivenKey extracts the key and value components of the "env" map's
+// combined "key=value" field, given the "key" component.
+//
+// This will work as long as the length of the supplied key equals the length of
+// entry's key prefix, regardless of the key's contents which means that
+// normalized keys can be used.
+//
+// It assumes that the environment variable is well-formed, one of:
+// - KEY (returns "")
+// - KEY= (returns "")
+// - KEY=VALUE (returns "VALUE")
+//
+// If this is not the case, the result is undefined. However, this is an
+// internal function and should never encounter a situation where the entry is
+// not well-formed and prefixed by the supplied key.
+func splitEntryGivenKey(key, entry string) (k, v string) {
+        prefixSize := len(key) + len("=")
+        switch {
+        case len(entry) < prefixSize:
+                return entry, ""
+        case len(entry) == prefixSize:
+                return entry[:len(key)], ""
+        default:
+                return entry[:len(key)], entry[prefixSize:]
+        }
+}