Merge pull request #178 from marineam/env

Write /etc/environment

initialize/config.go

@@ -223,12 +223,27 @@ func Apply(cfg CloudConfig, env *Environment) error {
 		cfg.Coreos.Units = append(cfg.Coreos.Units, u...)
 	}
 
+	wroteEnvironment := false
 	for _, file := range cfg.WriteFiles {
-		path, err := system.WriteFile(&file, env.Root())
+		fullPath, err := system.WriteFile(&file, env.Root())
 		if err != nil {
 			return err
 		}
-		log.Printf("Wrote file %s to filesystem", path)
+		if path.Clean(file.Path) == "/etc/environment" {
+			wroteEnvironment = true
+		}
+		log.Printf("Wrote file %s to filesystem", fullPath)
+	}
+
+	if !wroteEnvironment {
+		ef := env.DefaultEnvironmentFile()
+		if ef != nil {
+			err := system.WriteEnvFile(ef, env.Root())
+			if err != nil {
+				return err
+			}
+			log.Printf("Updated /etc/environment")
+		}
 	}
 
 	if env.NetconfType() != "" {

initialize/env.go

@@ -4,6 +4,8 @@ import (
 	"os"
 	"path"
 	"strings"
+
+	"github.com/coreos/coreos-cloudinit/system"
 )
 
 const DefaultSSHKeyName = "coreos-cloudinit"
@@ -65,6 +67,26 @@ func (e *Environment) Apply(data string) string {
 	return data
 }
 
+func (e *Environment) DefaultEnvironmentFile() *system.EnvFile {
+	ef := system.EnvFile{
+		File: &system.File{
+			Path: "/etc/environment",
+		},
+		Vars: map[string]string{},
+	}
+	if ip, ok := e.substitutions["$public_ipv4"]; ok && len(ip) > 0 {
+		ef.Vars["COREOS_PUBLIC_IPV4"] = ip
+	}
+	if ip, ok := e.substitutions["$private_ipv4"]; ok && len(ip) > 0 {
+		ef.Vars["COREOS_PRIVATE_IPV4"] = ip
+	}
+	if len(ef.Vars) == 0 {
+		return nil
+	} else {
+		return &ef
+	}
+}
+
 // normalizeSvcEnv standardizes the keys of the map (environment variables for a service)
 // by replacing any dashes with underscores and ensuring they are entirely upper case.
 // For example, "some-env" --> "SOME_ENV"

initialize/env_test.go

@@ -1,8 +1,12 @@
 package initialize
 
 import (
+	"io/ioutil"
 	"os"
+	"path"
 	"testing"
+
+	"github.com/coreos/coreos-cloudinit/system"
 )
 
 func TestEnvironmentApply(t *testing.T) {
@@ -56,3 +60,47 @@ ExecStop=/usr/bin/echo $unknown`,
 		}
 	}
 }
+
+func TestEnvironmentFile(t *testing.T) {
+	subs := map[string]string{
+		"$public_ipv4":  "1.2.3.4",
+		"$private_ipv4": "5.6.7.8",
+	}
+	expect := "COREOS_PUBLIC_IPV4=1.2.3.4\nCOREOS_PRIVATE_IPV4=5.6.7.8\n"
+
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	env := NewEnvironment("./", "./", "./", "", "", subs)
+	ef := env.DefaultEnvironmentFile()
+	err = system.WriteEnvFile(ef, dir)
+	if err != nil {
+		t.Fatalf("WriteEnvFile failed: %v", err)
+	}
+
+	fullPath := path.Join(dir, "etc", "environment")
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != expect {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+}
+
+func TestEnvironmentFileNil(t *testing.T) {
+	subs := map[string]string{
+		"$public_ipv4":  "",
+		"$private_ipv4": "",
+	}
+
+	env := NewEnvironment("./", "./", "./", "", "", subs)
+	ef := env.DefaultEnvironmentFile()
+	if ef != nil {
+		t.Fatalf("Environment file not nil: %v", ef)
+	}
+}

system/env_file.go

@@ -0,0 +1,89 @@
+package system
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path"
+	"regexp"
+)
+
+type EnvFile struct {
+	Vars map[string]string
+	// mask File.Content, it shouldn't be used.
+	Content interface{} `json:"-" yaml:"-"`
+	*File
+}
+
+// only allow sh compatible identifiers
+var validKey = regexp.MustCompile(`^[a-zA-Z0-9_]+$`)
+
+// match each line, optionally capturing valid identifiers, discarding dos line endings
+var lineLexer = regexp.MustCompile(`(?m)^((?:([a-zA-Z0-9_]+)=)?.*?)\r?\n`)
+
+// mergeEnvContents: Update the existing file contents with new values,
+// preserving variable ordering and all content this code doesn't understand.
+// All new values are appended to the bottom of the old.
+func mergeEnvContents(old []byte, pending map[string]string) []byte {
+	var buf bytes.Buffer
+	var match [][]byte
+
+	// it is awkward for the regex to handle a missing newline gracefully
+	if len(old) != 0 && !bytes.HasSuffix(old, []byte{'\n'}) {
+		old = append(old, byte('\n'))
+	}
+
+	for _, match = range lineLexer.FindAllSubmatch(old, -1) {
+		key := string(match[2])
+		if value, ok := pending[key]; ok {
+			fmt.Fprintf(&buf, "%s=%s\n", key, value)
+			delete(pending, key)
+		} else {
+			fmt.Fprintf(&buf, "%s\n", match[1])
+		}
+	}
+
+	for key, value := range pending {
+		fmt.Fprintf(&buf, "%s=%s\n", key, value)
+	}
+
+	return buf.Bytes()
+}
+
+// WriteEnvFile updates an existing env `KEY=value` formated file with
+// new values provided in EnvFile.Vars; File.Content is ignored.
+// Existing ordering and any unknown formatting such as comments are
+// preserved. If no changes are required the file is untouched.
+func WriteEnvFile(ef *EnvFile, root string) error {
+	// validate new keys, mergeEnvContents uses pending to track writes
+	pending := make(map[string]string, len(ef.Vars))
+	for key, value := range ef.Vars {
+		if !validKey.MatchString(key) {
+			return fmt.Errorf("Invalid name %q for %s", key, ef.Path)
+		}
+		pending[key] = value
+	}
+
+	if len(pending) == 0 {
+		return nil
+	}
+
+	oldContent, err := ioutil.ReadFile(path.Join(root, ef.Path))
+	if err != nil {
+		if os.IsNotExist(err) {
+			oldContent = []byte{}
+		} else {
+			return err
+		}
+	}
+
+	newContent := mergeEnvContents(oldContent, pending)
+	if bytes.Equal(oldContent, newContent) {
+		return nil
+	}
+
+	ef.File.Content = string(newContent)
+	_, err = WriteFile(ef.File, root)
+	return err
+}

system/env_file_test.go

@@ -0,0 +1,385 @@
+package system
+
+import (
+	"io/ioutil"
+	"os"
+	"path"
+	"strings"
+	"testing"
+)
+
+const (
+	base          = "# a file\nFOO=base\n\nBAR= hi there\n"
+	baseNoNewline = "# a file\nFOO=base\n\nBAR= hi there"
+	baseDos       = "# a file\r\nFOO=base\r\n\r\nBAR= hi there\r\n"
+	expectUpdate  = "# a file\nFOO=test\n\nBAR= hi there\nNEW=a value\n"
+	expectCreate  = "FOO=test\nNEW=a value\n"
+)
+
+var (
+	valueUpdate = map[string]string{
+		"FOO": "test",
+		"NEW": "a value",
+	}
+	valueNoop = map[string]string{
+		"FOO": "base",
+	}
+	valueEmpty   = map[string]string{}
+	valueInvalid = map[string]string{
+		"FOO-X": "test",
+	}
+)
+
+func TestWriteEnvFileUpdate(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(base), 0644)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueUpdate,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != expectUpdate {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+}
+
+func TestWriteEnvFileUpdateNoNewline(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(baseNoNewline), 0644)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueUpdate,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != expectUpdate {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+}
+
+func TestWriteEnvFileCreate(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueUpdate,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != expectCreate {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+}
+
+func TestWriteEnvFileNoop(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(base), 0644)
+
+	oldStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueNoop,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != base {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+
+	newStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	if oldStat.ModTime() != newStat.ModTime() {
+		t.Fatal("File mtime changed.")
+	}
+}
+
+func TestWriteEnvFileUpdateDos(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(baseDos), 0644)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueUpdate,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != expectUpdate {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+}
+
+// A middle ground noop, values are unchanged but we did have a value.
+// Seems reasonable to rewrite the file in Unix format anyway.
+func TestWriteEnvFileDos2Unix(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(baseDos), 0644)
+
+	oldStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueNoop,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != base {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+
+	newStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	if oldStat.ModTime() != newStat.ModTime() {
+		t.Fatal("File mtime changed.")
+	}
+}
+
+// If it really is a noop (structure is empty) don't even do dos2unix
+func TestWriteEnvFileEmpty(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(baseDos), 0644)
+
+	oldStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueEmpty,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err != nil {
+		t.Fatalf("Unable to read expected file: %v", err)
+	}
+
+	if string(contents) != baseDos {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	}
+
+	newStat, err := os.Stat(fullPath)
+	if err != nil {
+		t.Fatal("Unable to stat file: %v", err)
+	}
+
+	if oldStat.ModTime() != newStat.ModTime() {
+		t.Fatal("File mtime changed.")
+	}
+}
+
+// no point in creating empty files
+func TestWriteEnvFileEmptyNoCreate(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueEmpty,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err != nil {
+		t.Fatalf("WriteFile failed: %v", err)
+	}
+
+	contents, err := ioutil.ReadFile(fullPath)
+	if err == nil {
+		t.Fatalf("File has incorrect contents: %q", contents)
+	} else if !os.IsNotExist(err) {
+		t.Fatalf("Unexpected error while reading file: %v", err)
+	}
+}
+
+func TestWriteEnvFilePermFailure(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(base), 0000)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueUpdate,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if !os.IsPermission(err) {
+		t.Fatalf("Not a pemission denied error: %v", err)
+	}
+}
+
+func TestWriteEnvFileNameFailure(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+
+	name := "foo.conf"
+	fullPath := path.Join(dir, name)
+	ioutil.WriteFile(fullPath, []byte(base), 0000)
+
+	ef := EnvFile{
+		File: &File{
+			Path: name,
+		},
+		Vars: valueInvalid,
+	}
+
+	err = WriteEnvFile(&ef, dir)
+	if err == nil || !strings.HasPrefix(err.Error(), "Invalid name") {
+		t.Fatalf("Not an invalid name error: %v", err)
+	}
+}