From 08230a8fbc4acf5d57d1f3aebc57e6882b2844f9 Mon Sep 17 00:00:00 2001
From: Kason Braley <kbraley9@gmail.com>
Date: Sat, 6 Jul 2024 10:48:47 -0700
Subject: [PATCH] Format modified AST to buffer first before writing to file

This should avoid all potential scenarios where `format.Node` would fail
and the source file would be written to and end up empty. Now, we use a
buffer first to catch any of those errors. Only if it doesn't fail, we
open the source file and write the formatted modified AST to it.
---
 snap.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/snap.go b/snap.go
index 6d104dd..ed4a31b 100644
--- a/snap.go
+++ b/snap.go
@@ -54,6 +54,7 @@ import (
 	"go/format"
 	"go/parser"
 	"go/token"
+	"io"
 	"os"
 	"runtime"
 	"strings"
@@ -160,16 +161,24 @@ func (s *Snapshot) Diff(got string) {
 		return true
 	})
 
+	// Format the modified AST to a buffer first to avoid writing garbage(or nothing at all) back
+	// to the source file. Only if this succeeds, we then flush the buffer to the source file.
+	var buf bytes.Buffer
+	if err := format.Node(&buf, fset, f); err != nil {
+		s.t.Errorf("snap: Failed to format modified AST, aborting: %s", err)
+		return
+	}
+
 	outFile, err := os.OpenFile(s.location.file, os.O_WRONLY|os.O_TRUNC, 0644)
 	if err != nil {
-		s.t.Errorf("snap: Failed to open file for writing: %s", err)
+		s.t.Errorf("snap: Failed to open source file %q for writing to: %s", s.location.file, err)
 		return
 	}
 	defer outFile.Close()
 
-	// Write the modified AST back to the original file.
-	if err := format.Node(outFile, fset, f); err != nil {
-		s.t.Errorf("snap: Failed to write modified AST: %s", err)
+	// Write the modified(and formatted) AST in the buffer back to the original source file.
+	if _, err := io.Copy(outFile, &buf); err != nil {
+		s.t.Errorf("snap: Failed to write modified AST to source file: %s", err)
 		return
 	}