From c1edf992f994d531d84a25b7b59e09205419a045 Mon Sep 17 00:00:00 2001
From: Diff <diff@tilde.town>
Date: Wed, 24 Mar 2021 06:01:31 +0000
Subject: [PATCH] Make sure all requests are complete before shutting down.

Missed a spot in the net/http docs. Calling Server.Shutdown() causes
ListenAndServe to return *immediately*, and when it does main() exits and
the "Upload Success" message gets cut off before it can be fully sent, which
spooks HTTP clients a bit.

To avoid that, you have to wait for the call to Server.Shutdown() to return,
and then signal back to main() that it's now safe to exit.
---
 main.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index bf01e7b..03e4520 100644
--- a/main.go
+++ b/main.go
@@ -86,6 +86,8 @@ func getIP() (string, error) {
 	return localAddr.IP.String(), nil
 }
 
+var done = make(chan struct{})
+
 func main() {
 	conf, err := getConfig()
 	if err != nil {
@@ -122,6 +124,12 @@ func main() {
 		fmt.Println(err)
 		os.Exit(1)
 	}
+	
+	// Wait for the server to finish any transfers, up to 3 seconds
+	select {
+		case <-done:
+		case <-time.After(3*time.Second):
+	}
 }
 
 func setupDownload(server *http.Server, conf Config) {
@@ -133,7 +141,7 @@ func setupDownload(server *http.Server, conf Config) {
 		
 		downloads--
 		if downloads == 0 {
-			server.Shutdown(context.Background())
+			go shutdown(server)
 		}
 	})
 }
@@ -213,6 +221,11 @@ func setupUpload(server *http.Server, conf Config) {
 		}
 
 		tpl.ExecuteTemplate(w, "UploadMessage", "Upload successful!")
-		server.Shutdown(context.Background())
+		go shutdown(server)
 	})
 }
+
+func shutdown(server *http.Server) {
+	server.Shutdown(context.Background())
+	done <- struct{}{}
+}
\ No newline at end of file