system: new file writer for updating env-style files

This can be used to safely update config files cloudinit does not have
exclusive control over. For example update.conf or /etc/environment.

system/env_file.go

@@ -0,0 +1,86 @@
+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`)
+
+func updateEnv(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, updateEnv 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 := updateEnv(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)
+	}
+}