Feature: Added more documentation and clean code. (#20)

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

@@ -2,21 +2,33 @@ package main
 
 import (
 	"context"
+	"flag"
 	"log"
 	"time"
 	"waffle/internal/packet"
 
 	"waffle/internal/worker"
 )
 
-const networkInterfaceDescription = "Intel(R) I211 Gigabit Network Connection"
+const (
+	defaultNetworkInterfaceDescription = "Intel(R) I211 Gigabit Network Connection"
+)
 
 func main() {
+	var (
+		networkInterface string
+	)
+	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")
 
-	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
@@ -26,10 +38,10 @@ func main() {
 
 	collector := worker.NewCollector(
 		cfg,
-		packet.NewWindowsNetworkInterfaceProvider(networkInterfaceDescription),
-		inMemoryPacketSerializer)
+		packet.NewWindowsNetworkInterfaceProvider(networkInterface),
+		packetSerializer)
 
 	if err := collector.Run(ctx); err != nil {
-		panic(err.Error())
+		log.Fatalln("Error during running collector: ", err.Error())
 	}
 }

cmd/proxy/main.go

@@ -4,9 +4,8 @@ import (
 	"context"
 	"embed"
 	_ "embed"
-	"fmt"
-	"os"
-
+	"flag"
+	"log"
 	"waffle/cmd/proxy/server"
 )
 
@@ -17,9 +16,15 @@ var yamlConfigBytes []byte
 var certificates embed.FS
 
 func main() {
+	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")
+
 	ctx := context.Background()
-	if err := server.Run(ctx, yamlConfigBytes, certificates); err != nil {
-		_, _ = fmt.Fprintf(os.Stderr, "%s\n", err)
-		os.Exit(1)
+	if err := server.Run(ctx, proxyServerPort, visualizeServerPort, yamlConfigBytes, certificates); err != nil {
+		log.Fatalln(err)
 	}
 }

cmd/proxy/server/server.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"embed"
+	"fmt"
 	"log"
 	"os"
 	"os/signal"
@@ -19,7 +20,29 @@ import (
 	"waffle/internal/waf/guard"
 )
 
-func Run(ctx context.Context, yamlConfigBytes []byte, certificates embed.FS) error {
+// Run initializes and starts the Waffle Proxy server with the provided context, configuration, and embedded certificates.
+//
+// It first sets up signal handling to allow graceful shutdown on receiving an interrupt signal.
+//
+// The function loads environment-specific configurations, then parses the provided YAML configuration
+// to initialize a DNS provider for managing domain names.
+//
+// Next, it sets up the certificate provider using locally embedded certificates,
+// loading custom CA certificates, certificate PEM blocks, and key PEM blocks.
+//
+// A defense coordinator is initialized to handle security measures like XSS protection,
+// along with an in-memory rate limiter to control the number of requests allowed per a given time window (5 minutes).
+//
+// Additionally, a server for visualizing traffic is set up on port :8081.
+//
+// The WAF (Web Application Firewall) handler is constructed using a redirect handler, the defender (for security),
+// the rate limiter, and a visualizer from the visualization server.
+//
+// Finally, the main proxy server is started on port :8080, with the guard handler and certificate provider.
+// 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()
 
@@ -35,28 +58,47 @@ 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)
 
-	log.Println("Starting Waffle Proxy on port :8080 🚀")
+	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())
@@ -65,20 +107,33 @@ func Run(ctx context.Context, yamlConfigBytes []byte, certificates embed.FS) err
 	return nil
 }
 
-func loadLocalCustomCACerts(certificates embed.FS) [][]byte {
-	certBytes, _ := certificates.ReadFile(".cert/ca.crt")
-
-	return [][]byte{certBytes}
+// 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, error) {
+	certBytes, err := certificates.ReadFile(".cert/ca.crt")
+	if err != nil {
+		return nil, err
+	}
+	return [][]byte{certBytes}, nil
 }
 
-func loadLocalCertPEMBlock(certificates embed.FS) []byte {
-	certBytes, _ := certificates.ReadFile(".cert/server.crt")
-
-	return certBytes
+// 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, error) {
+	certBytes, err := certificates.ReadFile(".cert/server.crt")
+	if err != nil {
+		return nil, err
+	}
+	return certBytes, nil
 }
 
-func loadLocalKeyPEMBlock(certificates embed.FS) []byte {
-	certBytes, _ := certificates.ReadFile(".cert/server.key")
-
-	return certBytes
+// 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, 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)

internal/certificate/provider.go

@@ -13,12 +13,18 @@ type Provider interface {
 	GetCACertificatesPool() (*x509.CertPool, error)
 }
 
+// LocalCertificatesProvider is a structure that holds custom CA certificates,
+// the server certificate PEM block, and the private key PEM block.
+// It provides methods for retrieving TLS certificates and a CA certificate pool for establishing secure connections.
 type LocalCertificatesProvider struct {
 	customCaCerts [][]byte
 	certPEMBlock  []byte
 	keyPEMBlock   []byte
 }
 
+// NewLocalCertificatesProvider initializes a new instance of LocalCertificatesProvider.
+// It takes custom CA certificates, a certificate PEM block, and a key PEM block as input,
+// and returns a LocalCertificatesProvider that can provide certificates for secure connections.
 func NewLocalCertificatesProvider(caCerts [][]byte, certPEMBlock, keyPEMBlock []byte) *LocalCertificatesProvider {
 	return &LocalCertificatesProvider{
 		customCaCerts: caCerts,
@@ -27,6 +33,9 @@ func NewLocalCertificatesProvider(caCerts [][]byte, certPEMBlock, keyPEMBlock []
 	}
 }
 
+// GetTLSCertificate returns a TLS certificate created from the certificate PEM block and key PEM block.
+// This is used to provide the server's certificate for TLS/SSL handshakes.
+// If an error occurs while loading the key pair, it returns an error.
 func (l *LocalCertificatesProvider) GetTLSCertificate() (*tls.Certificate, error) {
 	cert, err := tls.X509KeyPair(l.certPEMBlock, l.keyPEMBlock)
 	if err != nil {
@@ -36,6 +45,10 @@ func (l *LocalCertificatesProvider) GetTLSCertificate() (*tls.Certificate, error
 	return &cert, nil
 }
 
+// GetCACertificatesPool returns a pool of CA certificates, including both system CA certificates and
+// any custom CA certificates provided to the LocalCertificatesProvider.
+// If the system CA certificate pool cannot be loaded, it returns an error.
+// Custom CA certificates are appended to the pool to support specific trusted CAs.
 func (l *LocalCertificatesProvider) GetCACertificatesPool() (*x509.CertPool, error) {
 	caCertPool, err := x509.SystemCertPool()
 	if err != nil {

internal/packet/find.go

@@ -11,6 +11,7 @@ import (
 )
 
 var (
+	// ErrNetworkInterfaceNotFound is an error which indicates  that network interface was not found
 	ErrNetworkInterfaceNotFound = errors.New("network interface not found")
 )
 
@@ -28,6 +29,9 @@ func NewWindowsNetworkInterfaceProvider(interfaceDescription string) *WindowsNet
 	return &WindowsNetworkInterfaceProvider{interfaceDescription: interfaceDescription}
 }
 
+// GetNetworkInterface retrieves all available interfaces, verifies if interface's description matches description
+// in interfaceDescription field of WindowsNetworkInterfaceProvider struct, and returns an interfaces which
+// suits that condition.
 func (w *WindowsNetworkInterfaceProvider) GetNetworkInterface() (*pcap.Interface, error) {
 	interfaces, err := pcap.FindAllDevs()
 	if err != nil {

internal/proxy/server.go

@@ -61,12 +61,18 @@ var (
 	alpnProto = "acme-tls/1"
 )
 
+// Server represents an HTTP server that listens on a specific address.
+// It uses a certificate provider to manage TLS certificates and an HTTP handler to process incoming requests.
 type Server struct {
 	addr                string
 	certificateProvider certificate.Provider
 	handler             http.Handler
 }
 
+// NewServer initializes a new Server instance with the given address, certificate provider, and HTTP handler.
+// The address defines where the server will listen for incoming connections.
+// The certificateProvider is used to manage the TLS certificates for secure connections.
+// The handler (usually a redirect handler) will process HTTP requests that are received by the server.
 func NewServer(
 	addr string,
 	certificateProvider certificate.Provider,
@@ -79,6 +85,22 @@ func NewServer(
 	}
 }
 
+// Start initializes and starts the server with TLS configuration using the provided certificate provider.
+// It first retrieves the CA certificates pool and server TLS certificate from the certificate provider.
+// Then, it configures the server to use TLS 1.3, sets the allowed cipher suites, and prepares the server for
+// mutual TLS authentication (client certificates are verified if provided).
+//
+// A TCP listener is created to listen for incoming TLS connections on the specified address.
+//
+// The HTTP server is configured with various timeouts (read, write, idle), a maximum header size, and a custom error logger.
+// The server handles incoming requests using the provided handler.
+//
+// A graceful shutdown mechanism is implemented: when an interrupt signal is received (e.g., Ctrl+C),
+// the server begins shutting down by closing the active listener and processing any outstanding requests.
+//
+// The function waits for all idle connections to be closed before returning.
+//
+// If there is any error while starting or serving the server, the error is returned.
 func (s *Server) Start() error {
 	caCertPool, err := s.certificateProvider.GetCACertificatesPool()
 	if err != nil {