1 files changed, 99 insertions, 26 deletions

diff --git a/src/cmd/gofmt/gofmt.go b/src/cmd/gofmt/gofmt.go
index bb22aea031..f4fb6bff84 100644
--- a/src/cmd/gofmt/gofmt.go
+++ b/src/cmd/gofmt/gofmt.go
@@ -17,10 +17,12 @@ import (
 	"internal/diff"
 	"io"
 	"io/fs"
+	"math/rand"
 	"os"
 	"path/filepath"
 	"runtime"
 	"runtime/pprof"
+	"strconv"
 	"strings"
 
 	"golang.org/x/sync/semaphore"
@@ -269,21 +271,9 @@ func processFile(filename string, info fs.FileInfo, in io.Reader, r *reporter) e
 			if info == nil {
 				panic("-w should not have been allowed with stdin")
 			}
-			// make a temporary backup before overwriting original
+
 			perm := info.Mode().Perm()
-			bakname, err := backupFile(filename+".", src, perm)
-			if err != nil {
-				return err
-			}
-			fdSem <- true
-			err = os.WriteFile(filename, res, perm)
-			<-fdSem
-			if err != nil {
-				os.Rename(bakname, filename)
-				return err
-			}
-			err = os.Remove(bakname)
-			if err != nil {
+			if err := writeFile(filename, src, res, perm, info.Size()); err != nil {
 				return err
 			}
 		}
@@ -467,28 +457,111 @@ func fileWeight(path string, info fs.FileInfo) int64 {
 	return info.Size()
 }
 
-// backupFile writes data to a new file named filename<number> with permissions perm,
-// with <number randomly chosen such that the file name is unique. backupFile returns
-// the chosen file name.
-func backupFile(filename string, data []byte, perm fs.FileMode) (string, error) {
+// writeFile updates a file with the new formatted data.
+func writeFile(filename string, orig, formatted []byte, perm fs.FileMode, size int64) error {
+	// Make a temporary backup file before rewriting the original file.
+	bakname, err := backupFile(filename, orig, perm)
+	if err != nil {
+		return err
+	}
+
 	fdSem <- true
 	defer func() { <-fdSem }()
 
-	// create backup file
-	f, err := os.CreateTemp(filepath.Dir(filename), filepath.Base(filename))
+	fout, err := os.OpenFile(filename, os.O_WRONLY, perm)
 	if err != nil {
-		return "", err
+		// We couldn't even open the file, so it should
+		// not have changed.
+		os.Remove(bakname)
+		return err
 	}
-	bakname := f.Name()
-	err = f.Chmod(perm)
+	defer fout.Close() // for error paths
+
+	restoreFail := func(err error) {
+		fmt.Fprintf(os.Stderr, "gofmt: %s: error restoring file to original: %v; backup in %s\n", filename, err, bakname)
+	}
+
+	n, err := fout.Write(formatted)
+	if err == nil && int64(n) < size {
+		err = fout.Truncate(int64(n))
+	}
+
 	if err != nil {
-		f.Close()
+		// Rewriting the file failed.
+
+		if n == 0 {
+			// Original file unchanged.
+			os.Remove(bakname)
+			return err
+		}
+
+		// Try to restore the original contents.
+
+		no, erro := fout.WriteAt(orig, 0)
+		if erro != nil {
+			// That failed too.
+			restoreFail(erro)
+			return err
+		}
+
+		if no < n {
+			// Original file is shorter. Truncate.
+			if erro = fout.Truncate(int64(no)); erro != nil {
+				restoreFail(erro)
+				return err
+			}
+		}
+
+		if erro := fout.Close(); erro != nil {
+			restoreFail(erro)
+			return err
+		}
+
+		// Original contents restored.
 		os.Remove(bakname)
-		return bakname, err
+		return err
+	}
+
+	if err := fout.Close(); err != nil {
+		restoreFail(err)
+		return err
+	}
+
+	// File updated.
+	os.Remove(bakname)
+	return nil
+}
+
+// backupFile writes data to a new file named filename<number> with permissions perm,
+// with <number> randomly chosen such that the file name is unique. backupFile returns
+// the chosen file name.
+func backupFile(filename string, data []byte, perm fs.FileMode) (string, error) {
+	fdSem <- true
+	defer func() { <-fdSem }()
+
+	nextRandom := func() string {
+		return strconv.Itoa(rand.Int())
+	}
+
+	dir, base := filepath.Split(filename)
+	var (
+		bakname string
+		f       *os.File
+	)
+	for {
+		bakname = filepath.Join(dir, base+"."+nextRandom())
+		var err error
+		f, err = os.OpenFile(bakname, os.O_RDWR|os.O_CREATE|os.O_EXCL, perm)
+		if err == nil {
+			break
+		}
+		if err != nil && !os.IsExist(err) {
+			return "", err
+		}
 	}
 
 	// write data to backup file
-	_, err = f.Write(data)
+	_, err := f.Write(data)
 	if err1 := f.Close(); err == nil {
 		err = err1
 	}