Merge branch 'main' of https://github.com/werniq/waffle into feature/k8s

Taskfile.yaml

@@ -0,0 +1,36 @@
+version: '3'
+
+tasks:
+  certs_windows:
+    cmds:
+      - mkdir -p "./cmd/proxy/.cert"
+      - openssl ecparam -out ./cmd/proxy/.cert/ca.key -name prime256v1 -genkey
+      - openssl req -new -sha256 -key ./cmd/proxy/.cert/ca.key -out ./cmd/proxy/.cert/ca.csr
+      - openssl x509 -req -sha256 -days 3650 -in ./cmd/proxy/.cert/ca.csr -signkey ./cmd/proxy/.cert/ca.key -out ./cmd/proxy/.cert/ca.crt
+      - openssl ecparam -out ./cmd/proxy/.cert/server.key -name prime256v1 -genkey
+      - openssl req -new -sha256 -key ./cmd/proxy/.cert/server.key -out ./cmd/proxy/.cert/server.csr
+      - openssl x509 -req -in ./cmd/proxy/.cert/server.csr -CA ./cmd/proxy/.cert/ca.crt -CAkey ./cmd/proxy/.cert/ca.key -CAcreateserial -out ./cmd/proxy/.cert/server.crt -days 3650 -sha256
+      - openssl x509 -in ./cmd/proxy/.cert/server.crt -text -noout
+
+  certs:
+    cmds:
+      - mkdir -p "./cmd/proxy/.cert"
+      - openssl ecparam -out ./cmd/proxy/.cert/ca.key -name prime256v1 -genkey
+      - openssl req -new -sha256 -key ./cmd/proxy/.cert/ca.key -out ./cmd/proxy/.cert/ca.csr
+      - openssl x509 -req -sha256 -days 3650 -in ./cmd/proxy/.cert/ca.csr -signkey ./cmd/proxy/.cert/ca.key -out ./cmd/proxy/.cert/ca.crt
+      - openssl ecparam -out ./cmd/proxy/.cert/server.key -name prime256v1 -genkey
+      - openssl req -new -sha256 -key ./cmd/proxy/.cert/server.key -out ./cmd/proxy/.cert/server.csr
+      - openssl x509 -req -in ./cmd/proxy/.cert/server.csr -CA ./cmd/proxy/.cert/ca.crt -CAkey ./cmd/proxy/.cert/ca.key -CAcreateserial -out ./cmd/proxy/.cert/server.crt -days 3650 -sha256
+      - openssl x509 -in ./cmd/proxy/.cert/server.crt -text -noout
+
+  mocks:
+    cmds:
+      - mockery
+
+  docker-build:
+    cmds:
+      - docker build -t ${IMAGE_NAME} -f .\build\Dockerfile .
+
+  docker-push:
+    cmds:
+      - docker push ${IMAGE_NAME}

build/Dockerfile

@@ -1,15 +1,16 @@
 FROM golang:1.22-alpine as builder
 
-ARG CGO_ENABLED=0
 WORKDIR /app
 
 COPY go.mod go.sum ./
 RUN go mod download
+
 COPY . .
 
-RUN go build cmd/proxy -o /out/app
+RUN go build -o /out/app ./cmd/proxy
 
-# multi stage build, to reduce size of the image
 FROM alpine:latest
+
 COPY --from=builder /out/app /app
+
 ENTRYPOINT ["/app"]

cmd/collector/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"flag"
 	"log"
 	"time"
@@ -20,11 +19,9 @@ func main() {
 	)
 	flag.StringVar(&networkInterface, "i", defaultNetworkInterfaceDescription, "Identification of the interface")
 
-	ctx := context.Background()
-
 	log.Println("starting collector")
 
-	inMemoryPacketSerializer := packet.NewMemoryPacketSerializer(time.Minute * 5)
+	packetSerializer := packet.NewMemoryPacketSerializer(time.Minute * 5)
 
 	// NEXT TODO: add BPF filter builder
 	// https://www.ibm.com/docs/en/qsip/7.4?topic=queries-berkeley-packet-filters
@@ -35,9 +32,9 @@ func main() {
 	collector := worker.NewCollector(
 		cfg,
 		packet.NewWindowsNetworkInterfaceProvider(networkInterface),
-		inMemoryPacketSerializer)
+		packetSerializer)
 
-	if err := collector.Run(ctx); err != nil {
-		panic(err.Error())
+	if err := collector.Run(); err != nil {
+		log.Fatalln("Error during running collector: ", err.Error())
 	}
 }

cmd/proxy/main.go

@@ -4,9 +4,12 @@ import (
 	"context"
 	"embed"
 	_ "embed"
-	"fmt"
+	"flag"
+	"log"
 	"os"
-
+	"os/signal"
+	"syscall"
+	"time"
 	"waffle/cmd/proxy/server"
 )
 
@@ -17,9 +20,32 @@ var yamlConfigBytes []byte
 var certificates embed.FS
 
 func main() {
-	ctx := context.Background()
-	if err := server.Run(ctx, yamlConfigBytes, certificates); err != nil {
-		_, _ = fmt.Fprintf(os.Stderr, "%s\n", err)
-		os.Exit(1)
+	var (
+		visualizeServerPort string
+		proxyServerPort     string
+	)
+	flag.StringVar(&visualizeServerPort, "p", "8081", "Port for server to listen on")
+	flag.StringVar(&proxyServerPort, "p", "8081", "Port for server to listen on")
+
+	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,9 +3,13 @@ package server
 import (
 	"context"
 	"embed"
+	"errors"
+	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"os/signal"
+	"syscall"
 	"time"
 	"waffle/internal/visualize"
 
@@ -41,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, 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())
@@ -57,54 +58,110 @@ func Run(ctx context.Context, yamlConfigBytes []byte, certificates embed.FS) err
 
 	yamlDnsProvider := domain.NewYamlNameSystemProvider(yamlCfg)
 
+	caCerts, err := loadLocalCustomCACerts(certificates)
+	if err != nil {
+		return err
+	}
+
+	certPemBlock, err := loadLocalCertPEMBlock(certificates)
+	if err != nil {
+		return err
+	}
+
+	keyPemBlock, err := loadLocalKeyPEMBlock(certificates)
+	if err != nil {
+		return err
+	}
+
 	certificateProvider := certificate.NewLocalCertificatesProvider(
-		loadLocalCustomCACerts(certificates),
-		loadLocalCertPEMBlock(certificates),
-		loadLocalKeyPEMBlock(certificates),
+		caCerts,
+		certPemBlock,
+		keyPemBlock,
 	)
 
 	defender := guard.NewDefenseCoordinator([]guard.Defender{&guard.XSS{}})
-
 	limiter := ratelimit.NewInMemoryLimiter(time.Minute * 5)
 
-	visualizeServer := visualize.NewServer(":8081")
+	visualizeServerPort = fmt.Sprintf(":%s", visualizeServerPort)
+
+	s := visualize.NewServer(visualizeServerPort)
 
 	guardHandler := waf.NewHandler(
 		redirect.NewHandler(yamlDnsProvider),
 		defender,
 		limiter,
-		visualizeServer.GetVisualizer(),
+		s.GetVisualizer(),
 	)
 
-	proxyServer := proxy.NewServer(":8080", certificateProvider, guardHandler)
+	proxyServerPort = fmt.Sprintf(":%s", proxyServerPort)
+	proxyServer := proxy.NewServer(proxyServerPort, certificateProvider, guardHandler)
 
-	log.Println("Starting Waffle Proxy on port :8080 🚀")
+	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.
-func loadLocalCustomCACerts(certificates embed.FS) [][]byte {
-	certBytes, _ := certificates.ReadFile(".cert/ca.crt")
-	return [][]byte{certBytes}
+func loadLocalCustomCACerts(certificates embed.FS) ([][]byte, error) {
+	certBytes, err := certificates.ReadFile(".cert/ca.crt")
+	if err != nil {
+		return nil, err
+	}
+	return [][]byte{certBytes}, nil
 }
 
 // loadLocalCertPEMBlock reads the local server certificate (server.crt) from the embedded file system.
 // It returns the certificate as a byte slice, which is later used to serve the server's public certificate in TLS/SSL connections.
-func loadLocalCertPEMBlock(certificates embed.FS) []byte {
-	certBytes, _ := certificates.ReadFile(".cert/server.crt")
-	return certBytes
+func loadLocalCertPEMBlock(certificates embed.FS) ([]byte, error) {
+	certBytes, err := certificates.ReadFile(".cert/server.crt")
+	if err != nil {
+		return nil, err
+	}
+	return certBytes, nil
 }
 
 // loadLocalKeyPEMBlock reads the private key (server.key) from the embedded file system.
 // It returns the private key as a byte slice, which is paired with the server certificate during TLS/SSL handshakes.
-func loadLocalKeyPEMBlock(certificates embed.FS) []byte {
-	certBytes, _ := certificates.ReadFile(".cert/server.key")
-	return certBytes
+func loadLocalKeyPEMBlock(certificates embed.FS) ([]byte, error) {
+	certBytes, err := certificates.ReadFile(".cert/server.key")
+	if err != nil {
+		return nil, err
+	}
+	return certBytes, nil
 }

cmd/tcpproxy/main.go

@@ -48,7 +48,7 @@ func tcpDummy() error {
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
-			return fmt.Errorf("liistener accept: %w", err)
+			return fmt.Errorf("listener accept: %w", err)
 		}
 
 		log.Println(conn)

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
 )