Safer file io for configuration files

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

2 files changed, 106 insertions, 0 deletions

diff --git a/pkg/ioutils/fswriters.go b/pkg/ioutils/fswriters.go
new file mode 100644
index 0000000000..ca97670724
--- /dev/null
+++ b/pkg/ioutils/fswriters.go
@@ -0,0 +1,75 @@
+package ioutils
+
+import (
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+)
+
+// NewAtomicFileWriter returns WriteCloser so that writing to it writes to a
+// temporary file and closing it atomically changes the temporary file to
+// destination path. Writing and closing concurrently is not allowed.
+func NewAtomicFileWriter(filename string, perm os.FileMode) (io.WriteCloser, error) {
+	f, err := ioutil.TempFile(filepath.Dir(filename), ".tmp-"+filepath.Base(filename))
+	if err != nil {
+		return nil, err
+	}
+	abspath, err := filepath.Abs(filename)
+	if err != nil {
+		return nil, err
+	}
+	return &atomicFileWriter{
+		f:  f,
+		fn: abspath,
+	}, nil
+}
+
+// AtomicWriteFile atomically writes data to a file named by filename.
+func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error {
+	f, err := NewAtomicFileWriter(filename, perm)
+	if err != nil {
+		return err
+	}
+	n, err := f.Write(data)
+	if err == nil && n < len(data) {
+		err = io.ErrShortWrite
+	}
+	if err1 := f.Close(); err == nil {
+		err = err1
+	}
+	return err
+}
+
+type atomicFileWriter struct {
+	f        *os.File
+	fn       string
+	writeErr error
+}
+
+func (w *atomicFileWriter) Write(dt []byte) (int, error) {
+	n, err := w.f.Write(dt)
+	if err != nil {
+		w.writeErr = err
+	}
+	return n, err
+}
+
+func (w *atomicFileWriter) Close() (retErr error) {
+	defer func() {
+		if retErr != nil {
+			os.Remove(w.f.Name())
+		}
+	}()
+	if err := w.f.Sync(); err != nil {
+		w.f.Close()
+		return err
+	}
+	if err := w.f.Close(); err != nil {
+		return err
+	}
+	if w.writeErr == nil {
+		return os.Rename(w.f.Name(), w.fn)
+	}
+	return nil
+}
diff --git a/pkg/ioutils/fswriters_test.go b/pkg/ioutils/fswriters_test.go
new file mode 100644
index 0000000000..40717a5108
--- /dev/null
+++ b/pkg/ioutils/fswriters_test.go
@@ -0,0 +1,31 @@
+package ioutils
+
+import (
+	"bytes"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestAtomicWriteToFile(t *testing.T) {
+	tmpDir, err := ioutil.TempDir("", "atomic-writers-test")
+	if err != nil {
+		t.Fatalf("Error when creating temporary directory: %s", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	expected := []byte("barbaz")
+	if err := AtomicWriteFile(filepath.Join(tmpDir, "foo"), expected, 0600); err != nil {
+		t.Fatalf("Error writing to file: %v", err)
+	}
+
+	actual, err := ioutil.ReadFile(filepath.Join(tmpDir, "foo"))
+	if err != nil {
+		t.Fatalf("Error reading from file: %v", err)
+	}
+
+	if bytes.Compare(actual, expected) != 0 {
+		t.Fatalf("Data mismatch, expected %q, got %q", expected, actual)
+	}
+}