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)