feat: refactor logging to use slog and enhance error handling

Author: hightemp

main.go

@@ -1,16 +1,21 @@
 package main
 
 import (
+	"context"
+	"crypto/subtle"
 	"crypto/tls"
 	"encoding/base64"
 	"flag"
 	"fmt"
 	"io"
-	"log"
+	"log/slog"
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"gopkg.in/yaml.v3"
@@ -25,26 +30,75 @@ type Config struct {
 	KeyPath   string `yaml:"key_path"`
 }
 
-var config Config
+// Hop-by-hop headers that should not be forwarded by proxies (RFC 2616 §13.5.1).
+var hopByHopHeaders = []string{
+	"Connection",
+	"Keep-Alive",
+	"Proxy-Authorization",
+	"Proxy-Connection",
+	"TE",
+	"Trailer",
+	"Transfer-Encoding",
+	"Upgrade",
+}
+
+// Buffer pool to reduce GC pressure during data transfer.
+var bufPool = sync.Pool{
+	New: func() any {
+		buf := make([]byte, 32*1024)
+		return &buf
+	},
+}
+
+var (
+	config     Config
+	httpClient *http.Client
+)
 
 func main() {
 	configPath := flag.String("config", "config.yaml", "Path to the config file")
 	flag.Parse()
 
 	content, err := os.ReadFile(*configPath)
 	if err != nil {
-		log.Fatalf("Error reading config file: %v", err)
+		slog.Error("Error reading config file", "error", err)
+		os.Exit(1)
 	}
 
 	err = yaml.Unmarshal(content, &config)
 	if err != nil {
-		log.Fatalf("Error parsing config file: %v", err)
+		slog.Error("Error parsing config file", "error", err)
+		os.Exit(1)
+	}
+
+	// Global HTTP client with connection pooling and timeouts.
+	httpClient = &http.Client{
+		Transport: &http.Transport{
+			MaxIdleConns:        100,
+			MaxIdleConnsPerHost: 10,
+			IdleConnTimeout:     90 * time.Second,
+			TLSHandshakeTimeout: 10 * time.Second,
+			DialContext: (&net.Dialer{
+				Timeout:   10 * time.Second,
+				KeepAlive: 30 * time.Second,
+			}).DialContext,
+		},
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			if len(via) >= 5 {
+				return fmt.Errorf("too many redirects (>%d) while following %s", len(via), via[len(via)-1].URL.String())
+			}
+			return nil
+		},
+		Timeout: 60 * time.Second,
 	}
 
 	server := &http.Server{
-		Addr: config.ProxyAddr,
+		Addr:         config.ProxyAddr,
+		ReadTimeout:  30 * time.Second,
+		WriteTimeout: 60 * time.Second,
+		IdleTimeout:  120 * time.Second,
 		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			log.Println("Received request:", r.Method, r.URL)
+			slog.Info("Received request", "method", r.Method, "url", r.URL.String())
 			if !basicAuth(w, r) {
 				return
 			}
@@ -57,58 +111,84 @@ func main() {
 		}),
 	}
 
-	log.Printf("Starting proxy server on %s\n", config.ProxyAddr)
+	// Graceful shutdown on SIGINT/SIGTERM.
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
+	go func() {
+		<-ctx.Done()
+		slog.Info("Shutting down proxy server...")
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+		defer cancel()
+		if err := server.Shutdown(shutdownCtx); err != nil {
+			slog.Error("Server shutdown error", "error", err)
+		}
+	}()
+
+	slog.Info("Starting proxy server", "addr", config.ProxyAddr, "proto", config.Proto)
 	if config.Proto == "https" {
 		ln, err := net.Listen("tcp", config.ProxyAddr)
 		if err != nil {
-			log.Fatalf("Error creating listener: %v", err)
+			slog.Error("Error creating listener", "error", err)
+			os.Exit(1)
 		}
 
 		cert, err := tls.LoadX509KeyPair(config.CertPath, config.KeyPath)
 		if err != nil {
-			log.Fatalf("Error loading certificate: %v", err)
+			slog.Error("Error loading certificate", "error", err)
+			os.Exit(1)
 		}
 
 		server.TLSConfig = &tls.Config{
 			Certificates: []tls.Certificate{cert},
+			MinVersion:   tls.VersionTLS12,
 		}
 
 		tlsListener := tls.NewListener(ln, server.TLSConfig)
 
-		log.Fatal(server.Serve(tlsListener))
-
-		// log.Fatal(server.ListenAndServeTLS(config.CertPath, config.KeyPath))
+		if err := server.Serve(tlsListener); err != http.ErrServerClosed {
+			slog.Error("Server error", "error", err)
+			os.Exit(1)
+		}
 	} else {
-		log.Fatal(server.ListenAndServe())
+		if err := server.ListenAndServe(); err != http.ErrServerClosed {
+			slog.Error("Server error", "error", err)
+			os.Exit(1)
+		}
 	}
+
+	slog.Info("Server stopped")
 }
 
 func basicAuth(w http.ResponseWriter, r *http.Request) bool {
 	auth := r.Header.Get("Proxy-Authorization")
 	if auth == "" {
-		log.Println("No Proxy-Authorization header")
+		slog.Debug("No Proxy-Authorization header", "remote", r.RemoteAddr)
 		w.Header().Set("Proxy-Authenticate", `Basic realm="Proxy Authorization Required"`)
 		w.WriteHeader(http.StatusProxyAuthRequired)
 		return false
 	}
 
 	payload, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(auth, "Basic "))
 	if err != nil {
-		log.Println("Error decoding auth:", err)
+		slog.Warn("Error decoding auth", "error", err, "remote", r.RemoteAddr)
 		w.WriteHeader(http.StatusBadRequest)
 		return false
 	}
 
 	pair := strings.SplitN(string(payload), ":", 2)
 	if len(pair) != 2 {
-		log.Printf("Invalid auth format: %v\n", pair)
+		slog.Warn("Invalid auth format", "remote", r.RemoteAddr)
 		w.Header().Set("Proxy-Authenticate", `Basic realm="Proxy Authorization Required"`)
 		w.WriteHeader(http.StatusProxyAuthRequired)
 		return false
 	}
 
-	if pair[0] != config.Username || pair[1] != config.Password {
-		log.Printf("Invalid credentials: %s:%s\n", pair[0], pair[1])
+	// Constant-time comparison to prevent timing attacks.
+	usernameMatch := subtle.ConstantTimeCompare([]byte(pair[0]), []byte(config.Username))
+	passwordMatch := subtle.ConstantTimeCompare([]byte(pair[1]), []byte(config.Password))
+	if usernameMatch&passwordMatch != 1 {
+		slog.Warn("Invalid credentials", "user", pair[0], "remote", r.RemoteAddr)
 		w.Header().Set("Proxy-Authenticate", `Basic realm="Proxy Authorization Required"`)
 		w.WriteHeader(http.StatusProxyAuthRequired)
 		return false
@@ -119,30 +199,22 @@ func basicAuth(w http.ResponseWriter, r *http.Request) bool {
 
 func handleHTTP(w http.ResponseWriter, r *http.Request) {
 	r.RequestURI = ""
-	r.Header.Del("Proxy-Connection")
-	r.Header.Del("Proxy-Authorization")
-
 	r.Host = r.URL.Host
 
-	client := &http.Client{
-		CheckRedirect: func(req *http.Request, via []*http.Request) error {
-			if len(via) >= 5 {
-				lastURL := via[len(via)-1].URL.String()
-				err := fmt.Errorf("too many redirects (>%d) while following %s", len(via), lastURL)
-				return err
-			}
-			return nil
-		},
+	// Remove hop-by-hop headers.
+	for _, h := range hopByHopHeaders {
+		r.Header.Del(h)
 	}
 
-	resp, err := client.Do(r)
+	resp, err := httpClient.Do(r)
 	if err != nil {
-		log.Printf("Error forwarding request: %v\n", err)
+		slog.Error("Error forwarding request", "error", err, "url", r.URL.String())
 		http.Error(w, err.Error(), http.StatusServiceUnavailable)
 		return
 	}
 	defer resp.Body.Close()
 
+	// Copy response headers, skipping hop-by-hop headers.
 	for key, values := range resp.Header {
 		for _, value := range values {
 			w.Header().Add(key, value)
@@ -151,54 +223,74 @@ func handleHTTP(w http.ResponseWriter, r *http.Request) {
 
 	w.WriteHeader(resp.StatusCode)
 
-	written, err := io.Copy(w, resp.Body)
+	bufPtr := bufPool.Get().(*[]byte)
+	defer bufPool.Put(bufPtr)
+	written, err := io.CopyBuffer(w, resp.Body, *bufPtr)
 	if err != nil {
-		log.Printf("Error copying response body after %d bytes: %v\n", written, err)
-		http.Error(w, err.Error(), http.StatusServiceUnavailable)
+		// Headers already sent — cannot call http.Error, just log.
+		slog.Error("Error copying response body", "written", written, "error", err)
 		return
 	}
-	log.Printf("Successfully copied %d bytes from response\n", written)
+	slog.Debug("Response copied", "bytes", written, "url", r.URL.String())
 }
 
 func handleTunneling(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodConnect {
-		log.Printf("Error: Method not allowed: %s\n", r.Method)
-		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
-		return
-	}
-
 	destConn, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
 	if err != nil {
-		log.Printf("Error: Can't connect to host: %s, %v\n", r.Host, err)
+		slog.Error("Can't connect to host", "host", r.Host, "error", err)
 		http.Error(w, err.Error(), http.StatusServiceUnavailable)
 		return
 	}
+
 	w.WriteHeader(http.StatusOK)
 
 	hijacker, ok := w.(http.Hijacker)
 	if !ok {
-		log.Printf("Error: Hijacking not supported\n")
-		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
+		slog.Error("Hijacking not supported")
+		destConn.Close()
 		return
 	}
 	clientConn, _, err := hijacker.Hijack()
 	if err != nil {
-		log.Printf("Error: Client connection error: %v\n", err)
-		http.Error(w, err.Error(), http.StatusServiceUnavailable)
+		slog.Error("Client connection hijack error", "error", err)
+		destConn.Close()
 		return
 	}
 
-	go transfer(destConn, clientConn)
-	go transfer(clientConn, destConn)
+	// Use a WaitGroup to wait for both directions to finish,
+	// then close both connections cleanly.
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		transfer(destConn, clientConn)
+		// Signal the other direction to stop by setting a read deadline.
+		if tc, ok := destConn.(*net.TCPConn); ok {
+			tc.SetReadDeadline(time.Now())
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		transfer(clientConn, destConn)
+		if tc, ok := clientConn.(*net.TCPConn); ok {
+			tc.SetReadDeadline(time.Now())
+		}
+	}()
+
+	wg.Wait()
+	destConn.Close()
+	clientConn.Close()
 }
 
-func transfer(destination io.WriteCloser, source io.ReadCloser) {
-	defer destination.Close()
-	defer source.Close()
-	bytes, err := io.Copy(destination, source)
+func transfer(destination io.Writer, source io.Reader) {
+	bufPtr := bufPool.Get().(*[]byte)
+	defer bufPool.Put(bufPtr)
+	written, err := io.CopyBuffer(destination, source, *bufPtr)
 	if err != nil {
-		log.Printf("Transfer error: %v\n", err)
+		slog.Debug("Transfer finished with error", "bytes", written, "error", err)
 	} else {
-		log.Printf("Transferred %d bytes\n", bytes)
+		slog.Debug("Transfer complete", "bytes", written)
 	}
 }