From 43bab58f0a588b22b97d040c83131e1c1bb44cee Mon Sep 17 00:00:00 2001
From: Christoph Haas <christoph.h@sprinternet.at>
Date: Mon, 22 Feb 2021 22:25:08 +0100
Subject: [PATCH] WIP: context for clean shutdown

---
 cmd/wg-portal/main.go   | 30 ++++++++++++++++++++++++++++--
 internal/server/core.go |  4 +++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/cmd/wg-portal/main.go b/cmd/wg-portal/main.go
index 10c2243..7febe38 100644
--- a/cmd/wg-portal/main.go
+++ b/cmd/wg-portal/main.go
@@ -1,8 +1,12 @@
 package main
 
 import (
+	"context"
 	"io/ioutil"
 	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/h44z/wg-portal/internal/server"
 	"github.com/sirupsen/logrus"
@@ -11,16 +15,38 @@ import (
 func main() {
 	_ = setupLogger(logrus.StandardLogger())
 
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+
 	logrus.Infof("Starting WireGuard Portal Server...")
 
+	// Context for clean shutdown
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	service := server.Server{}
-	if err := service.Setup(); err != nil {
+	if err := service.Setup(ctx); err != nil {
 		logrus.Fatalf("Setup failed: %v", err)
 	}
 
-	service.Run()
+	// Attach signal handlers to context
+	go func() {
+		osCall := <-c
+		logrus.Tracef("received system call: %v", osCall)
+		cancel() // cancel the context
+	}()
+
+	// Start main process in background
+	go service.Run()
+
+	<-ctx.Done() // Wait until the context gets canceled
+
+	// Give goroutines some time to stop gracefully
+	logrus.Info("Stopping WireGuard Portal Server...")
+	time.Sleep(2 * time.Second)
 
 	logrus.Infof("Stopped WireGuard Portal Server...")
+	logrus.Exit(0)
 }
 
 func setupLogger(logger *logrus.Logger) error {
diff --git a/internal/server/core.go b/internal/server/core.go
index 36aa368..3e743c9 100644
--- a/internal/server/core.go
+++ b/internal/server/core.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"encoding/gob"
 	"errors"
 	"html/template"
@@ -64,6 +65,7 @@ type StaticData struct {
 
 type Server struct {
 	// Core components
+	ctx     context.Context
 	config  *common.Config
 	server  *gin.Engine
 	users   *UserManager
@@ -79,7 +81,7 @@ type Server struct {
 	ldapCacheUpdater *ldap.UserCache
 }
 
-func (s *Server) Setup() error {
+func (s *Server) Setup(ctx context.Context) error {
 	dir := s.getExecutableDirectory()
 	rDir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
 	logrus.Infof("Real working directory: %s", rDir)