From cc0d76e016e6ff0a7acbeca64e47fe57eeb99c5f Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Sun, 12 Apr 2020 14:12:19 +0200
Subject: [PATCH] s2c: Handle sigint better (#239)

* s2c/s2d: Handle sigint.
* Tweak output

* Simplify.
---
 s2/cmd/s2c/main.go | 25 +++++++++++++++++++------
 s2/cmd/s2d/main.go | 17 ++++++++++++-----
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/s2/cmd/s2c/main.go b/s2/cmd/s2c/main.go
index 36969b63c5..a630de6893 100644
--- a/s2/cmd/s2c/main.go
+++ b/s2/cmd/s2c/main.go
@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"runtime"
 	"runtime/pprof"
@@ -79,10 +80,13 @@ Options:`)
 
 	// No args, use stdin/stdout
 	if len(args) == 1 && args[0] == "-" {
+		// Catch interrupt, so we don't exit at once.
+		// os.Stdin will return EOF, so we should be able to get everything.
+		signal.Notify(make(chan os.Signal), os.Interrupt)
 		wr.Reset(os.Stdout)
-		_, err := io.Copy(wr, os.Stdin)
-		exitErr(err)
-		exitErr(wr.Close())
+		_, err = wr.ReadFrom(os.Stdin)
+		printErr(err)
+		printErr(wr.Close())
 		return
 	}
 	var files []string
@@ -137,7 +141,7 @@ Options:`)
 				dstFilename = "(discarded)"
 			}
 			if !*quiet {
-				fmt.Println("Compressing", filename, "->", dstFilename)
+				fmt.Print("Compressing ", filename, " -> ", dstFilename)
 			}
 			// Input file.
 			file, err := os.Open(filename)
@@ -181,11 +185,14 @@ Options:`)
 				elapsed := time.Since(start)
 				mbpersec := (float64(input) / (1024 * 1024)) / (float64(elapsed) / (float64(time.Second)))
 				pct := float64(wc.n) * 100 / float64(input)
-				fmt.Printf("%d -> %d [%.02f%%]; %.01fMB/s\n", input, wc.n, pct, mbpersec)
+				fmt.Printf(" %d -> %d [%.02f%%]; %.01fMB/s\n", input, wc.n, pct, mbpersec)
 			}
 			if *remove {
 				closeOnce.Do(func() {
 					file.Close()
+					if !*quiet {
+						fmt.Println("Removing", filename)
+					}
 					err := os.Remove(filename)
 					exitErr(err)
 				})
@@ -194,9 +201,15 @@ Options:`)
 	}
 }
 
+func printErr(err error) {
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "\nERROR:", err.Error())
+	}
+}
+
 func exitErr(err error) {
 	if err != nil {
-		fmt.Fprintln(os.Stderr, "ERROR:", err.Error())
+		fmt.Fprintln(os.Stderr, "\nERROR:", err.Error())
 		os.Exit(2)
 	}
 }
diff --git a/s2/cmd/s2d/main.go b/s2/cmd/s2d/main.go
index ce72fde630..2fddbd7a86 100644
--- a/s2/cmd/s2d/main.go
+++ b/s2/cmd/s2d/main.go
@@ -19,6 +19,7 @@ import (
 
 var (
 	safe   = flag.Bool("safe", false, "Do not overwrite output files")
+	verify = flag.Bool("verify", false, "Verify files, but do not write output")
 	stdout = flag.Bool("c", false, "Write all output to stdout. Multiple input files will be concatenated")
 	remove = flag.Bool("rm", false, "Delete source file(s) after successful decompression")
 	quiet  = flag.Bool("q", false, "Don't write any output to terminal, except errors")
@@ -88,11 +89,14 @@ Options:`)
 		if *bench > 0 {
 			dstFilename = "(discarded)"
 		}
+		if *verify {
+			dstFilename = "(verify)"
+		}
 
 		func() {
 			var closeOnce sync.Once
 			if !*quiet {
-				fmt.Println("Decompressing", filename, "->", dstFilename)
+				fmt.Print("Decompressing ", filename, " -> ", dstFilename)
 			}
 			// Input file.
 			file, err := os.Open(filename)
@@ -113,7 +117,7 @@ Options:`)
 			}
 			var out io.Writer
 			switch {
-			case *bench > 0:
+			case *bench > 0 || *verify:
 				out = ioutil.Discard
 			case *stdout:
 				out = os.Stdout
@@ -133,11 +137,14 @@ Options:`)
 				elapsed := time.Since(start)
 				mbPerSec := (float64(output) / (1024 * 1024)) / (float64(elapsed) / (float64(time.Second)))
 				pct := float64(output) * 100 / float64(rc.n)
-				fmt.Printf("%d -> %d [%.02f%%]; %.01fMB/s\n", rc.n, output, pct, mbPerSec)
+				fmt.Printf(" %d -> %d [%.02f%%]; %.01fMB/s\n", rc.n, output, pct, mbPerSec)
 			}
-			if *remove {
+			if *remove && !*verify {
 				closeOnce.Do(func() {
 					file.Close()
+					if !*quiet {
+						fmt.Println("Removing", filename)
+					}
 					err := os.Remove(filename)
 					exitErr(err)
 				})
@@ -148,7 +155,7 @@ Options:`)
 
 func exitErr(err error) {
 	if err != nil {
-		fmt.Fprintln(os.Stderr, "ERROR:", err.Error())
+		fmt.Fprintln(os.Stderr, "\nERROR:", err.Error())
 		os.Exit(2)
 	}
 }