From 79c0a17328fd4da360621ef28ab620e148365339 Mon Sep 17 00:00:00 2001
From: Alberto Bertogli <albertito@blitiri.com.ar>
Date: Wed, 25 Jan 2017 11:04:29 +0000
Subject: [PATCH] safeio: Extend WriteFile to take operations before the final
 rename

This patch extends WriteFile to allow arbitrary operations to be applied
to the file before it is atomically renamed.

This will be used in upcoming patches to change the mtime of the file
before it is atomically renamed.
---
 internal/safeio/safeio.go      | 17 ++++++++-
 internal/safeio/safeio_test.go | 65 ++++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 3 deletions(-)

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).