From 997eb2eb2c92c331d0299bbcf9a257903a2b3f0c Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Wed, 25 Oct 2023 14:45:13 +0200
Subject: [PATCH] cmd/geth: add gz support, improve error messages

---
 cmd/geth/chaincmd.go |  6 +++---
 cmd/utils/cmd.go     | 18 +++++++++++++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/cmd/geth/chaincmd.go b/cmd/geth/chaincmd.go
index fec7830644..80e4a271f9 100644
--- a/cmd/geth/chaincmd.go
+++ b/cmd/geth/chaincmd.go
@@ -419,7 +419,7 @@ func exportPreimages(ctx *cli.Context) error {
 
 // exportSnapshotPreimages dumps the preimage data to a flat file.
 func exportSnapshotPreimages(ctx *cli.Context) error {
-	if ctx.Args().Len() < 1 {
+	if ctx.NArg() < 1 {
 		utils.Fatalf("This command requires an argument.")
 	}
 	stack, _ := makeConfigNode(ctx)
@@ -428,10 +428,10 @@ func exportSnapshotPreimages(ctx *cli.Context) error {
 	chain, _ := utils.MakeChain(ctx, stack, true)
 
 	var root common.Hash
-	if ctx.Args().Len() > 1 {
+	if ctx.NArg() > 1 {
 		rootBytes := common.FromHex(ctx.Args().Get(1))
 		if len(rootBytes) != common.HashLength {
-			return fmt.Errorf("invalid root hash length")
+			return fmt.Errorf("invalid hash: %s", ctx.Args().Get(1))
 		}
 		root = common.BytesToHash(rootBytes)
 	}
diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 5a377685e5..cfd8dfac8f 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -385,8 +385,17 @@ func ExportSnapshotPreimages(chain *core.BlockChain, fn string, root common.Hash
 	}
 	defer fh.Close()
 
-	writer := bufio.NewWriter(fh)
-	defer writer.Flush()
+	var writer io.Writer = fh
+
+	if strings.HasSuffix(fn, ".gz") {
+		gz := gzip.NewWriter(writer)
+		defer gz.Close()
+		writer = gz
+	}
+
+	buf := bufio.NewWriter(writer)
+	defer buf.Flush()
+	writer = buf
 
 	statedb, err := chain.State()
 	if err != nil {
@@ -441,8 +450,11 @@ func ExportSnapshotPreimages(chain *core.BlockChain, fn string, root common.Hash
 
 	for item := range hashCh {
 		preimage := rawdb.ReadPreimage(statedb.Database().DiskDB(), item.Hash)
+		if len(preimage) == 0 {
+			return fmt.Errorf("missing preimage for %v", item.Hash)
+		}
 		if len(preimage) != item.Size {
-			return fmt.Errorf("invalid preimage size")
+			return fmt.Errorf("invalid preimage size, have %d", len(preimage))
 		}
 		if _, err := writer.Write(preimage); err != nil {
 			return fmt.Errorf("failed to write preimage: %w", err)