Implemented graceful shutdown in server/server.go

cmd/proxy/server/server.go

@@ -3,8 +3,13 @@ package server
 import (
 	"context"
 	"embed"
+	"errors"
 	"fmt"
 	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 	"waffle/internal/visualize"
 
@@ -75,11 +80,9 @@ func Run(proxyServerPort, visualizeServerPort string, yamlConfigBytes []byte, ce
 	)
 
 	defender := guard.NewDefenseCoordinator([]guard.Defender{&guard.XSS{}})
-
 	limiter := ratelimit.NewInMemoryLimiter(time.Minute * 5)
 
 	visualizeServerPort = fmt.Sprintf(":%s", visualizeServerPort)
-
 	s := visualize.NewServer(visualizeServerPort)
 
 	guardHandler := waf.NewHandler(
@@ -90,13 +93,37 @@ func Run(proxyServerPort, visualizeServerPort string, yamlConfigBytes []byte, ce
 	)
 
 	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

internal/proxy/server.go

@@ -166,3 +166,9 @@ func (s *Server) Start() error {
 
 	return nil
 }
+
+// Shutdown is a function which will be called when gracefully shutting down the server.
+func (s *Server) Shutdown(ctx context.Context) error {
+	log.Println("Server shut down gracefully")
+	return nil
+}

internal/worker/collector.go

@@ -67,7 +67,7 @@ func NewCollector(
 //
 // Deferred actions ensure that the packet channel is closed and a log message is generated
 // when the collector is closed
-func (c *Collector) Run(ctx context.Context) error {
+func (c *Collector) Run() error {
 	netInterface, err := c.netInterfaceProvider.GetNetworkInterface()
 	if err != nil {
 		return fmt.Errorf("get network interface using net interface provider: %w", err)
@@ -89,12 +89,12 @@ func (c *Collector) Run(ctx context.Context) error {
 	signalChan := make(chan os.Signal, 1)
 	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
 
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	defer func() {
 		close(packetsChan)
-		log.Println("collector closed")
+		log.Println("Collector closed")
 	}()
 
 	go func() {
@@ -109,7 +109,6 @@ func (c *Collector) Run(ctx context.Context) error {
 			log.Println("Received shutdown signal")
 			cancel()
 		case <-ctx.Done():
-
 		}
 	}()