Cleaner filemode options

README.md

@@ -1,6 +1,8 @@
 # atomic
-    import "github.com/natefinch/atomic"
-atomic is a go package for atomic file writing
+
+Go package for atomic file writing
+
+[![Go Reference](https://pkg.go.dev/badge/github.com/natefinch/atomic.svg)](https://pkg.go.dev/github.com/natefinch/atomic)
 
 By default, writing to a file in go (and generally any language) can fail
 partway through... you then have a partially written file, which probably was
@@ -25,11 +27,30 @@ change either file.
 
 ## func WriteFile
 ``` go
-func WriteFile(filename string, r io.Reader) (err error)
+func WriteFile(filename string, r io.Reader, opts ...Option) (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).  Additional option arguments can be used to change the
+default configuration for the target file.
+
+
+## Example
 
+``` go
+import (
+	"strings"
+
+	"github.com/natefinch/atomic"
+)
+
+func main() {
+	r := strings.NewReader("yes\n")
+	err := atomic.WriteFile("consistent.txt", r, atomic.FileMode(0440))
+	if err != nil {
+		// handle error
+	}
+}
+```

atomic.go

@@ -9,12 +9,55 @@ import (
 	"path/filepath"
 )
 
+// FileOptions define the behaviour of `FileWrite()`.
+type FileOptions struct {
+	defaultFileMode os.FileMode
+	fileMode        os.FileMode
+	keepFileMode    bool
+}
+
+// Option functions modify FileOptions.
+type Option func(*FileOptions)
+
+// FileMode sets the file mode to the desired value and has precedence over all
+// other options.
+func FileMode(mode os.FileMode) Option {
+	return func(opts *FileOptions) {
+		opts.fileMode = mode
+	}
+}
+
+// DefaultFileMode sets the default file mode instead of using the
+// `ioutil.TempFile()` default of `0600`.
+func DefaultFileMode(mode os.FileMode) Option {
+	return func(opts *FileOptions) {
+		opts.defaultFileMode = mode
+	}
+}
+
+// KeepFileMode preserves the file mode of an existing file instead of using the
+// default value.
+func KeepFileMode(keep bool) Option {
+	return func(opts *FileOptions) {
+		opts.keepFileMode = keep
+	}
+}
+
 // 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).
-func WriteFile(filename string, r io.Reader) (err error) {
+// is unmodified).  Additional option arguments can be used to change the
+// default configuration for the target file.
+func WriteFile(filename string, r io.Reader, opts ...Option) (err error) {
+	// original behaviour is to preserve the mode of an existing file.
+	fopts := &FileOptions{
+		keepFileMode: true,
+	}
+	for _, opt := range opts {
+		opt(fopts)
+	}
+
 	// write to a temp file first, then we'll atomically replace the target file
 	// with the temp file.
 	dir, file := filepath.Split(filename)
@@ -43,27 +86,39 @@ func WriteFile(filename string, r io.Reader) (err error) {
 	if err := f.Sync(); err != nil {
 		return fmt.Errorf("can't flush tempfile %q: %v", name, err)
 	}
+	// get file info via file descriptor before closing it.
+	sourceInfo, err := f.Stat()
+	if err != nil {
+		return err
+	}
 	if err := f.Close(); err != nil {
 		return fmt.Errorf("can't close tempfile %q: %v", name, err)
 	}
 
+	var fileMode os.FileMode
+	// change default file mode for when file does not exist yet.
+	if fopts.defaultFileMode != 0 {
+		fileMode = fopts.defaultFileMode
+	}
 	// get the file mode from the original file and use that for the replacement
 	// file, too.
-	destInfo, err := os.Stat(filename)
-	if os.IsNotExist(err) {
-		// no original file
-	} else if err != nil {
-		return err
-	} else {
-		sourceInfo, err := os.Stat(name)
-		if err != nil {
+	if fopts.keepFileMode {
+		destInfo, err := os.Stat(filename)
+		if err != nil && !os.IsNotExist(err) {
 			return err
 		}
-
-		if sourceInfo.Mode() != destInfo.Mode() {
-			if err := os.Chmod(name, destInfo.Mode()); err != nil {
-				return fmt.Errorf("can't set filemode on tempfile %q: %v", name, err)
-			}
+		if destInfo != nil {
+			fileMode = destInfo.Mode()
+		}
+	}
+	// given file mode always takes precedence
+	if fopts.fileMode != 0 {
+		fileMode = fopts.fileMode
+	}
+	// apply possible file mode change
+	if fileMode != 0 && fileMode != sourceInfo.Mode() {
+		if err := os.Chmod(name, fileMode); err != nil {
+			return fmt.Errorf("can't set filemode on tempfile %q: %v", name, err)
 		}
 	}
 	if err := ReplaceFile(name, filename); err != nil {

atomic_test.go

@@ -0,0 +1,83 @@
+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 TestWriteDefaultFileMode(t *testing.T) {
+	file := "bar.txt"
+	content := bytes.NewBufferString("bar")
+	defer func() { _ = os.Remove(file) }()
+	if err := WriteFile(file, content, DefaultFileMode(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: %v", fi.Mode())
+	}
+	// check if file mode is preserved
+	if err := os.Chmod(file, 0600); err != nil {
+		t.Errorf("Failed to change file mode: %q: %v", file, err)
+	}
+	if err := WriteFile(file, content, DefaultFileMode(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() != 0600 {
+		t.Errorf("File mode not correct: %v", fi.Mode())
+	}
+}
+
+func TestWriteFileMode(t *testing.T) {
+	file := "baz.txt"
+	content := bytes.NewBufferString("baz")
+	defer func() { _ = os.Remove(file) }()
+	if err := WriteFile(file, content, FileMode(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: %v", fi.Mode())
+	}
+	// ensure previous file mode is ingored
+	if err := os.Chmod(file, 0600); err != nil {
+		t.Errorf("Failed to change file mode: %q: %v", file, err)
+	}
+	if err := WriteFile(file, content, FileMode(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: %v", fi.Mode())
+	}
+}