Implement support for optional file mode (fixes natefinch/atomic#7)

README.md

@@ -25,11 +25,13 @@ change either file.
 
 ## func WriteFile
 ``` go
-func WriteFile(filename string, r io.Reader) (err error)
+func WriteFile(filename string, r io.Reader, mode ...os.FileMode) (err error)
 ```
 WriteFile atomically writes the contents of r to the specified filepath.  If
 an error occurs, the target file is guaranteed to be either fully written, or
 not written at all.  WriteFile overwrites any file that exists at the
 location (but only if the write fully succeeds, otherwise the existing file
-is unmodified).
+is unmodified). Permissions are copied from an existing file or the optional
+default file mode that can be given to be used instead of the default `0600`
+from ioutil.TempFile().
 

atomic.go

@@ -13,8 +13,10 @@ import (
 // an error occurs, the target file is guaranteed to be either fully written, or
 // not written at all.  WriteFile overwrites any file that exists at the
 // location (but only if the write fully succeeds, otherwise the existing file
-// is unmodified).
-func WriteFile(filename string, r io.Reader) (err error) {
+// is unmodified). Permissions are copied from an existing file or the optional
+// default file mode that can be given to be used instead of the default `0600`
+// from ioutil.TempFile().
+func WriteFile(filename string, r io.Reader, mode ...os.FileMode) (err error) {
 	// write to a temp file first, then we'll atomically replace the target file
 	// with the temp file.
 	dir, file := filepath.Split(filename)
@@ -39,6 +41,12 @@ func WriteFile(filename string, r io.Reader) (err error) {
 	if _, err := io.Copy(f, r); err != nil {
 		return fmt.Errorf("cannot write data to tempfile %q: %v", name, err)
 	}
+	// when optional mode was given change default mode of temp file.
+	if len(mode) > 0 {
+		if err := f.Chmod(mode[0]); err != nil {
+			return fmt.Errorf("cannot change file mode %q: %v", name, err);
+		}
+	}
 	if err := f.Close(); err != nil {
 		return fmt.Errorf("can't close tempfile %q: %v", name, err)
 	}

atomic_test.go

@@ -0,0 +1,39 @@
+package atomic
+
+import (
+	"bytes"
+	"os"
+	"testing"
+)
+
+func TestWriteFile(t *testing.T) {
+	file := "foo.txt"
+	content := bytes.NewBufferString("foo")
+	defer func() { _ = os.Remove(file) }()
+	if err := WriteFile(file, content); err != nil {
+		t.Errorf("Failed to write file: %q: %v", file, err)
+	}
+	fi, err := os.Stat(file)
+	if err != nil {
+		t.Errorf("Failed to stat file: %q: %v", file, err)
+	}
+	if fi.Mode() != 0600 {
+		t.Errorf("File mode not correct")
+	}
+}
+
+func TestWriteFileMode(t *testing.T) {
+	file := "bar.txt"
+	content := bytes.NewBufferString("bar")
+	defer func() { _ = os.Remove(file) }()
+	if err := WriteFile(file, content, 0644); err != nil {
+		t.Errorf("Failed to write file: %q: %v", file, err)
+	}
+	fi, err := os.Stat(file)
+	if err != nil {
+		t.Errorf("Failed to stat file: %q: %v", file, err)
+	}
+	if fi.Mode() != 0644 {
+		t.Errorf("File mode not correct")
+	}
+}