git » chasquid » commit 79c0a17

diff --git a/internal/safeio/safeio.go b/internal/safeio/safeio.go
index 8eba2a6..a78302d 100644
--- a/internal/safeio/safeio.go
+++ b/internal/safeio/safeio.go
@@ -9,13 +9,21 @@ import (
 	"syscall"
 )
 
+// Type FileOp represents an operation on a file (passed by its name).
+type FileOp func(fname string) error
+
 // WriteFile writes data to a file named by filename, atomically.
+//
 // It's a wrapper to ioutil.WriteFile, but provides atomicity (and increased
 // safety) by writing to a temporary file and renaming it at the end.
 //
+// Before the final rename, the given ops (if any) are called. They can be
+// used to manipulate the file before it is atomically renamed.
+// If any operation fails, the file is removed and the error is returned.
+//
 // Note this relies on same-directory Rename being atomic, which holds in most
 // reasonably modern filesystems.
-func WriteFile(filename string, data []byte, perm os.FileMode) error {
+func WriteFile(filename string, data []byte, perm os.FileMode, ops ...FileOp) error {
 	// Note we create the temporary file in the same directory, otherwise we
 	// would have no expectation of Rename being atomic.
 	// We make the file names start with "." so there's no confusion with the
@@ -50,6 +58,13 @@ func WriteFile(filename string, data []byte, perm os.FileMode) error {
 		return err
 	}
 
+	for _, op := range ops {
+		if err = op(tmpf.Name()); err != nil {
+			os.Remove(tmpf.Name())
+			return err
+		}
+	}
+
 	return os.Rename(tmpf.Name(), filename)
 }
 
diff --git a/internal/safeio/safeio_test.go b/internal/safeio/safeio_test.go
index 13f3a2e..f5dc81d 100644
--- a/internal/safeio/safeio_test.go
+++ b/internal/safeio/safeio_test.go
@@ -2,9 +2,11 @@ package safeio
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strings"
 	"testing"
 )
 
@@ -24,8 +26,8 @@ func mustTempDir(t *testing.T) string {
 	return dir
 }
 
-func testWriteFile(fname string, data []byte, perm os.FileMode) error {
-	err := WriteFile("file1", data, perm)
+func testWriteFile(fname string, data []byte, perm os.FileMode, ops ...FileOp) error {
+	err := WriteFile("file1", data, perm, ops...)
 	if err != nil {
 		return fmt.Errorf("error writing new file: %v", err)
 	}
@@ -81,6 +83,65 @@ func TestWriteFile(t *testing.T) {
 	}
 }
 
+func TestWriteFileWithOp(t *testing.T) {
+	dir := mustTempDir(t)
+
+	var opFile string
+	op := func(f string) error {
+		opFile = f
+		return nil
+	}
+
+	content := []byte("content 1")
+	if err := testWriteFile("file1", content, 0660, op); err != nil {
+		t.Error(err)
+	}
+
+	if opFile == "" {
+		t.Error("operation was not called")
+	}
+	if !strings.Contains(opFile, "file1") {
+		t.Errorf("operation called with suspicious file: %s", opFile)
+	}
+
+	// Remove the test directory, but only if we have not failed. We want to
+	// keep the failed structure for debugging.
+	if !t.Failed() {
+		os.RemoveAll(dir)
+	}
+}
+
+func TestWriteFileWithFailingOp(t *testing.T) {
+	dir := mustTempDir(t)
+
+	var opFile string
+	opOK := func(f string) error {
+		opFile = f
+		return nil
+	}
+
+	opError := errors.New("operation failed")
+	opFail := func(f string) error {
+		return opError
+	}
+
+	content := []byte("content 1")
+	err := WriteFile("file1", content, 0660, opOK, opOK, opFail)
+	if err != opError {
+		t.Errorf("different error, got %v, expected %v", err, opError)
+	}
+
+	if _, err := os.Stat(opFile); err == nil {
+		t.Errorf("temporary file was not removed after failure (%v)", opFile)
+	}
+
+	// Remove the test directory, but only if we have not failed. We want to
+	// keep the failed structure for debugging.
+	if !t.Failed() {
+		os.RemoveAll(dir)
+	}
+}
+
 // TODO: We should test the possible failure scenarios for WriteFile, but it
 // gets tricky without being able to do failure injection (or turning the code
 // into a mess).