From 81824be3bff2ce1572a49a10abe654aae02ecac0 Mon Sep 17 00:00:00 2001 From: Michael Marineau Date: Thu, 10 Jul 2014 14:42:28 -0700 Subject: [PATCH] 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 | 86 +++++++++ system/env_file_test.go | 385 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 471 insertions(+) create mode 100644 system/env_file.go create mode 100644 system/env_file_test.go diff --git a/system/env_file.go b/system/env_file.go new file mode 100644 index 0000000..b860d48 --- /dev/null +++ b/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 +} diff --git a/system/env_file_test.go b/system/env_file_test.go new file mode 100644 index 0000000..c3fc0a4 --- /dev/null +++ b/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) + } +}