From 9d5ee15ff86760beb3daaa1962eb556e60bbd107 Mon Sep 17 00:00:00 2001
From: Nick White <git@njw.name>
Date: Tue, 13 Oct 2020 16:24:45 +0100
Subject: [iiifdownloader] Catch SIGINT when writing a file to remove
 half-written files before exit

---
 cmd/iiifdownloader/main.go | 43 +++++++++++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 10 deletions(-)

(limited to 'cmd/iiifdownloader/main.go')

diff --git a/cmd/iiifdownloader/main.go b/cmd/iiifdownloader/main.go
index eeb253e..20c7ac8 100644
--- a/cmd/iiifdownloader/main.go
+++ b/cmd/iiifdownloader/main.go
@@ -11,9 +11,11 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"os/signal"
 	"path"
 	"strconv"
 	"strings"
+	"syscall"
 )
 
 // TODO: Add tests
@@ -238,19 +240,40 @@ func dlPage(bookdir, u string) error {
 		return fmt.Errorf("Error downloading page %s: %v", u, err)
 	}
 
-	f, err := os.Create(fn)
-	defer f.Close()
-	if err != nil {
-		return fmt.Errorf("Error creating file %s: %v\n", fn, err)
-	}
-	_, err = io.Copy(f, resp.Body)
-	if err != nil {
-		_ = f.Close()
+	// We do the file writing in a goroutine so that we can catch
+	// SIGINT (Ctrl-C) and ensure the file is removed, to ensure it
+	// can't be left in a half-written state. This is important so
+	// SIGINT can be used to stop the download in a state that it
+	// can safely be continued later by rerunning iiifdownloader.
+
+	sigint := make(chan os.Signal)
+	done := make(chan error)
+	signal.Notify(sigint, syscall.SIGINT)
+
+	go func() {
+		f, err := os.Create(fn)
+		defer f.Close()
+		if err != nil {
+			done <- fmt.Errorf("Error creating file %s: %v\n", fn, err)
+		}
+		_, err = io.Copy(f, resp.Body)
+		if err != nil {
+			_ = f.Close()
+			_ = os.Remove(fn)
+			done <- fmt.Errorf("Error writing file %s: %v\n", fn, err)
+		}
+		done <- nil
+	}()
+
+	select {
+	case <-sigint:
 		_ = os.Remove(fn)
-		return fmt.Errorf("Error writing file %s: %v\n", fn, err)
+		os.Exit(0)
+	case err = <-done:
 	}
 
-	return nil
+	signal.Reset(syscall.SIGINT)
+	return err
 }
 
 // dlNoPgNums downloads all pages, starting from zero, until either
-- 
cgit v1.2.1-24-ge1ad