cipd: accept a function in EnsureFile

client/cipd/local/fs.go

Index: client/cipd/local/fs.go
diff --git a/client/cipd/local/fs.go b/client/cipd/local/fs.go
index da4043e386c1c7faecf34e6900e67c33b454d4a8..80bffe13dd72dbf99c8878024b97e7c3ced52241 100644
--- a/client/cipd/local/fs.go
+++ b/client/cipd/local/fs.go
@@ -44,9 +44,9 @@ type FileSystem interface {
 	// directory path to the symlink if necessary.
 	EnsureSymlink(path string, target string) error
 
-	// EnsureFile creates a file with given content. If will create full directory
-	// path to the file if necessary.
-	EnsureFile(path string, body io.Reader) error
+	// EnsureFile creates a file and calls the function to write file content.
+	// It will create full directory path to the file if necessary.
+	EnsureFile(path string, write func(*os.File) error) error
 
 	// EnsureFileGone removes a file, logging the errors (if any). Missing file is
 	// not an error.
@@ -76,20 +76,29 @@ func NewFileSystem(root string, logger logging.Logger) FileSystem {
 	return &fsImpl{abs, logger}
 }
 
+// EnsureFile creates a file with the given content.
+// It will create full directory path to the file if necessary.
+func EnsureFile(fs FileSystem, path string, content io.Reader) error {
+	return fs.EnsureFile(path, func(f *os.File) error {
+		_, err := io.Copy(f, content)
+		return err
+	})
+}
+
 // fsImplErr implements FileSystem by returning given error from all methods.
 type fsImplErr struct {
 	err error
 }
 
-func (f *fsImplErr) Root() string                                   { return "" }
-func (f *fsImplErr) CwdRelToAbs(path string) (string, error)        { return "", f.err }
-func (f *fsImplErr) RootRelToAbs(path string) (string, error)       { return "", f.err }
-func (f *fsImplErr) EnsureDirectory(path string) (string, error)    { return "", f.err }
-func (f *fsImplErr) EnsureSymlink(path string, target string) error { return f.err }
-func (f *fsImplErr) EnsureFile(path string, body io.Reader) error   { return f.err }
-func (f *fsImplErr) EnsureFileGone(path string) error               { return f.err }
-func (f *fsImplErr) EnsureDirectoryGone(path string) error          { return f.err }
-func (f *fsImplErr) Replace(oldpath, newpath string) error          { return f.err }
+func (f *fsImplErr) Root() string                                             { return "" }
+func (f *fsImplErr) CwdRelToAbs(path string) (string, error)                  { return "", f.err }
+func (f *fsImplErr) RootRelToAbs(path string) (string, error)                 { return "", f.err }
+func (f *fsImplErr) EnsureDirectory(path string) (string, error)              { return "", f.err }
+func (f *fsImplErr) EnsureSymlink(path string, target string) error           { return f.err }
+func (f *fsImplErr) EnsureFile(path string, write func(*os.File) error) error { return f.err }
+func (f *fsImplErr) EnsureFileGone(path string) error                         { return f.err }
+func (f *fsImplErr) EnsureDirectoryGone(path string) error                    { return f.err }
+func (f *fsImplErr) Replace(oldpath, newpath string) error                    { return f.err }
 
 /// Implementation.
 
@@ -138,7 +147,7 @@ func (f *fsImpl) EnsureDirectory(path string) (string, error) {
 	return path, nil
 }
 
-func (f *fsImpl) EnsureFile(path string, body io.Reader) error {
+func (f *fsImpl) EnsureFile(path string, write func(*os.File) error) error {
 	path, err := f.CwdRelToAbs(path)
 	if err != nil {
 		return err
@@ -149,7 +158,7 @@ func (f *fsImpl) EnsureFile(path string, body io.Reader) error {
 
 	// Create a temp file with new content.
 	temp := tempFileName(path)
-	if err := createFile(temp, body); err != nil {
+	if err := createFile(temp, write); err != nil {
 		return err
 	}
 
@@ -316,8 +325,8 @@ func pseudoRand() string {
 	return fmt.Sprintf("%v_%v", os.Getpid(), ts)
 }
 
-// createFile creates a file with the given content.
-func createFile(path string, content io.Reader) (err error) {
+// createFile creates a file and calls the function to write file content.
+func createFile(path string, write func(*os.File) error) (err error) {
 	file, err := os.Create(path)
 	if err != nil {
 		return err
@@ -328,6 +337,5 @@ func createFile(path string, content io.Reader) (err error) {
 			err = closeErr
 		}
 	}()
-	_, err = io.Copy(file, content)
-	return err
+	return write(file)
 }