From 39bf3ad7f123119b61b2204cd7c97ea4964e67ad Mon Sep 17 00:00:00 2001
From: Harvey Tindall <hrfee@protonmail.ch>
Date: Wed, 5 Aug 2020 16:58:24 +0100
Subject: [PATCH] Safe shutdown

---
 api.go  | 10 +++++-----
 main.go | 25 ++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/api.go b/api.go
index 20fdcd7..07f051d 100644
--- a/api.go
+++ b/api.go
@@ -7,7 +7,7 @@ import (
 	"github.com/lithammer/shortuuid/v3"
 	"gopkg.in/ini.v1"
 	"os"
-	//"os/exec"
+	"os/signal"
 	"syscall"
 	"time"
 )
@@ -583,7 +583,7 @@ func (ctx *appContext) ModifyConfig(gc *gin.Context) {
 	gc.JSON(200, map[string]bool{"success": true})
 	if req["restart-program"].(bool) {
 		ctx.info.Println("Restarting...")
-		err := Restart()
+		err := ctx.Restart()
 		if err != nil {
 			ctx.err.Printf("Couldn't restart, try restarting manually. (%s)", err)
 		}
@@ -634,10 +634,11 @@ func (ctx *appContext) ModifyConfig(gc *gin.Context) {
 // 	panic(fmt.Errorf("restarting"))
 // }
 
-func Restart() error {
+func (ctx *appContext) Restart() error {
 	defer func() {
 		if r := recover(); r != nil {
-			os.Exit(0)
+			signal.Notify(ctx.quit, os.Interrupt)
+			<-ctx.quit
 		}
 	}()
 	args := os.Args
@@ -648,7 +649,6 @@ func Restart() error {
 		os.Setenv("JFA_EXEC", args[0])
 	}
 	env := os.Environ()
-	fmt.Printf("EXECUTABLE: %s\n", os.Getenv("JFA_EXEC"))
 	err := syscall.Exec(os.Getenv("JFA_EXEC"), []string{""}, env)
 	if err != nil {
 		return err
diff --git a/main.go b/main.go
index 21cedf4..b31d38d 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/base64"
 	"encoding/json"
@@ -14,7 +15,9 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"time"
 )
@@ -48,6 +51,7 @@ type appContext struct {
 	host             string
 	port             int
 	version          string
+	quit             chan os.Signal
 }
 
 func GenerateSecret(length int) (string, error) {
@@ -319,5 +323,24 @@ func main() {
 		router.POST("/modifyConfig", ctx.ModifyConfig)
 		ctx.info.Printf("Loading setup @ %s", address)
 	}
-	router.Run(address)
+	// router.Run(address)
+	srv := &http.Server{
+		Addr:    address,
+		Handler: router,
+	}
+	go func() {
+		if err := srv.ListenAndServe(); err != nil {
+			ctx.err.Printf("Failure serving: %s", err)
+		}
+	}()
+	ctx.quit = make(chan os.Signal)
+	signal.Notify(ctx.quit, os.Interrupt)
+	<-ctx.quit
+	ctx.info.Println("Shutting down...")
+
+	cntx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+	if err := srv.Shutdown(cntx); err != nil {
+		ctx.err.Fatalf("Server shutdown error: %s", err)
+	}
 }