Feature: Implement Graceful shutdown

cmd/collector/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"flag"
 	"log"
 	"time"
@@ -20,12 +19,6 @@ func main() {
 	)
 	flag.StringVar(&networkInterface, "i", defaultNetworkInterfaceDescription, "Identification of the interface")
 
-	// question: Why do we need context here? It is not used in collector.Run, except of ctx.Done, but since it is not
-	// context.WithTimeout (as example) it can not be closed in any way.
-	// Same in c.serializer.SerializePackets(ctx, packetsChan), it can not be closed there as well.
-	// Why not just to remove it?
-	ctx := context.Background()
-
 	log.Println("starting collector")
 
 	packetSerializer := packet.NewMemoryPacketSerializer(time.Minute * 5)
@@ -41,7 +34,7 @@ func main() {
 		packet.NewWindowsNetworkInterfaceProvider(networkInterface),
 		packetSerializer)
 
-	if err := collector.Run(ctx); err != nil {
+	if err := collector.Run(); err != nil {
 		log.Fatalln("Error during running collector: ", err.Error())
 	}
 }

cmd/proxy/main.go

@@ -6,6 +6,10 @@ import (
 	_ "embed"
 	"flag"
 	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
 	"waffle/cmd/proxy/server"
 )
 
@@ -23,8 +27,25 @@ func main() {
 	flag.StringVar(&visualizeServerPort, "p", "8081", "Port for server to listen on")
 	flag.StringVar(&proxyServerPort, "p", "8081", "Port for server to listen on")
 
-	ctx := context.Background()
-	if err := server.Run(ctx, proxyServerPort, visualizeServerPort, yamlConfigBytes, certificates); err != nil {
-		log.Fatalln(err)
+	quit := make(chan os.Signal)
+
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		if err := server.Run(proxyServerPort, visualizeServerPort, yamlConfigBytes, certificates); err != nil {
+			log.Fatalln(err)
+		}
+	}()
+
+	log.Println("Server started on :8080")
+
+	<-quit
+	log.Println("Shutdown signal received, shutting down gracefully...")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	if err := server.Shutdown(ctx); err != nil {
+		log.Fatalf("Server forced to shutdown: %s\n", err)
 	}
 }

cmd/proxy/server/server.go

@@ -3,10 +3,13 @@ package server
 import (
 	"context"
 	"embed"
+	"errors"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"os/signal"
+	"syscall"
 	"time"
 	"waffle/internal/visualize"
 
@@ -42,10 +45,7 @@ import (
 // If the proxy server fails to start, the function logs a fatal error.
 //
 // The function returns nil upon normal completion.
-func Run(ctx context.Context, proxyServerPort, visualizeServerPort string, yamlConfigBytes []byte, certificates embed.FS) error {
-	_, cancel := signal.NotifyContext(ctx, os.Interrupt)
-	defer cancel()
-
+func Run(proxyServerPort, visualizeServerPort string, yamlConfigBytes []byte, certificates embed.FS) error {
 	_, err := config.LoadEnvironmentConfig()
 	if err != nil {
 		log.Fatal(err.Error())
@@ -80,7 +80,6 @@ func Run(ctx context.Context, proxyServerPort, visualizeServerPort string, yamlC
 	)
 
 	defender := guard.NewDefenseCoordinator([]guard.Defender{&guard.XSS{}})
-
 	limiter := ratelimit.NewInMemoryLimiter(time.Minute * 5)
 
 	visualizeServerPort = fmt.Sprintf(":%s", visualizeServerPort)
@@ -95,18 +94,47 @@ func Run(ctx context.Context, proxyServerPort, visualizeServerPort string, yamlC
 	)
 
 	proxyServerPort = fmt.Sprintf(":%s", proxyServerPort)
-
 	proxyServer := proxy.NewServer(proxyServerPort, certificateProvider, guardHandler)
 
 	log.Printf("Starting Waffle Proxy on port %s 🚀\n", proxyServerPort)
 
-	if err := proxyServer.Start(); err != nil {
-		log.Fatal(err.Error())
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		if err := proxyServer.Start(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatalf("Proxy server encountered an error: %v\n", err)
+		}
+	}()
+
+	select {
+	case sig := <-signalChan:
+		log.Printf("Received shutdown signal: %s, shutting down gracefully...", sig)
+		cancel()
+	case <-ctx.Done():
+		log.Println("Context canceled, shutting down...")
+	}
+
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer shutdownCancel()
+
+	if err := proxyServer.Shutdown(shutdownCtx); err != nil {
+		log.Printf("Error during server shutdown: %v", err)
+	} else {
+		log.Println("Proxy server shut down gracefully.")
 	}
 
 	return nil
 }
 
+// Shutdown is a function to be called while 'gracefully shutdowning' the server.
+func Shutdown(ctx context.Context) error {
+	return nil
+}
+
 // loadLocalCustomCACerts reads the local custom CA certificates from the embedded file system.
 // It reads the CA certificate file (ca.crt) located in the ".cert" directory and returns it as a slice of byte slices.
 // This CA certificate is used for establishing trust during TLS/SSL handshakes.

go.mod

@@ -4,29 +4,31 @@ go 1.22.0
 
 require (
 	github.com/Netflix/go-env v0.0.0-20220526054621-78278af1949d
+	github.com/cdipaolo/goml v0.0.0-20220715001353-00e0c845ae1c
 	github.com/corazawaf/libinjection-go v0.1.3
 	github.com/emirpasic/gods v1.18.1
 	github.com/goccy/go-yaml v1.11.3
+	github.com/google/gopacket v1.1.19
 	github.com/google/uuid v1.6.0
+	github.com/malaschitz/randomForest v0.0.0-20240228214944-c64ffe1648e0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/stretchr/testify v1.9.0
+	golang.org/x/sys v0.17.0
 	nhooyr.io/websocket v1.8.10
 )
 
 require (
-	github.com/cdipaolo/goml v0.0.0-20220715001353-00e0c845ae1c // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fatih/color v1.16.0 // indirect
-	github.com/google/gopacket v1.1.19 // indirect
-	github.com/malaschitz/randomForest v0.0.0-20240228214944-c64ffe1648e0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/sjwhitworth/golearn v0.0.0-20221228163002-74ae077eafb2 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
 	golang.org/x/crypto v0.14.0 // indirect
-	golang.org/x/sys v0.17.0 // indirect
+	golang.org/x/net v0.10.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
 	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
 	gonum.org/v1/gonum v0.9.3 // indirect
+	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )