Merge pull request #57 from coder/yevhenii/proxy-tests

refactor: add tests for proxy package

Author: evgeniy-scherbina

boundary.go

@@ -62,12 +62,11 @@ func (b *Boundary) Start() error {
 	}
 
 	// Start proxy server in background
-	go func() {
-		err := b.proxyServer.Start(b.ctx)
-		if err != nil {
-			b.logger.Error("Proxy server error", "error", err)
-		}
-	}()
+	err = b.proxyServer.Start()
+	if err != nil {
+		b.logger.Error("Proxy server error", "error", err)
+		return err
+	}
 
 	// Give proxy time to start
 	time.Sleep(100 * time.Millisecond)
@@ -90,4 +89,4 @@ func (b *Boundary) Close() error {
 
 	// Close jailer
 	return b.jailer.Close()
-}
+}

go.mod

@@ -2,12 +2,16 @@ module github.com/coder/boundary
 
 go 1.24
 
-require github.com/coder/serpent v0.10.0
+require (
+	github.com/coder/serpent v0.10.0
+	github.com/stretchr/testify v1.8.4
+)
 
 require (
 	cdr.dev/slog v1.6.2-0.20240126064726-20367d4aede6 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/coder/pretty v0.0.0-20230908205945-e89ba86370e0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
@@ -18,6 +22,7 @@ require (
 	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/pion/transport/v2 v2.0.0 // indirect
 	github.com/pion/udp v0.1.4 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/rivo/uniseg v0.4.4 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	go.opentelemetry.io/otel v1.19.0 // indirect

proxy/proxy.go

@@ -2,8 +2,8 @@ package proxy
 
 import (
 	"bufio"
-	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -12,7 +12,7 @@ import (
 	"net/url"
 	"strings"
 	"sync"
-	"time"
+	"sync/atomic"
 
 	"github.com/coder/boundary/audit"
 	"github.com/coder/boundary/rules"
@@ -25,8 +25,9 @@ type Server struct {
 	logger     *slog.Logger
 	tlsConfig  *tls.Config
 	httpPort   int
+	started    atomic.Bool
 
-	httpServer *http.Server
+	listener net.Listener
 }
 
 // Config holds configuration for the proxy server
@@ -50,64 +51,70 @@ func NewProxyServer(config Config) *Server {
 }
 
 // Start starts the HTTP proxy server with TLS termination capability
-func (p *Server) Start(ctx context.Context) error {
-	// Create HTTP server with TLS termination capability
-	p.httpServer = &http.Server{
-		Addr:    fmt.Sprintf(":%d", p.httpPort),
-		Handler: http.HandlerFunc(p.handleHTTPWithTLSTermination),
+func (p *Server) Start() error {
+	if p.isStarted() {
+		return nil
+	}
+
+	p.logger.Info("Starting HTTP proxy with TLS termination", "port", p.httpPort)
+	var err error
+	p.listener, err = net.Listen("tcp", fmt.Sprintf(":%d", p.httpPort))
+	if err != nil {
+		p.logger.Error("Failed to create HTTP listener", "error", err)
+		return err
 	}
 
+	p.started.Store(true)
+
 	// Start HTTP server with custom listener for TLS detection
 	go func() {
-		p.logger.Info("Starting HTTP proxy with TLS termination", "port", p.httpPort)
-		listener, err := net.Listen("tcp", fmt.Sprintf(":%d", p.httpPort))
-		if err != nil {
-			p.logger.Error("Failed to create HTTP listener", "error", err)
-			return
-		}
-
 		for {
-			conn, err := listener.Accept()
+			conn, err := p.listener.Accept()
+			if err != nil && errors.Is(err, net.ErrClosed) && p.isStopped() {
+				return
+			}
 			if err != nil {
-				select {
-				case <-ctx.Done():
-					err = listener.Close()
-					if err != nil {
-						p.logger.Error("Failed to close listener", "error", err)
-					}
-					return
-				default:
-					p.logger.Error("Failed to accept connection", "error", err)
-					continue
-				}
+				p.logger.Error("Failed to accept connection", "error", err)
+				continue
 			}
 
 			// Handle connection with TLS detection
 			go p.handleConnectionWithTLSDetection(conn)
 		}
 	}()
 
-	// Wait for context cancellation
-	<-ctx.Done()
-	return p.Stop()
+	return nil
 }
 
 // Stops proxy server
 func (p *Server) Stop() error {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
+	if p.isStopped() {
+		return nil
+	}
+	p.started.Store(false)
 
-	var httpErr error
-	if p.httpServer != nil {
-		httpErr = p.httpServer.Shutdown(ctx)
+	if p.listener == nil {
+		p.logger.Error("unexpected nil listener")
+		return errors.New("unexpected nil listener")
 	}
 
-	if httpErr != nil {
-		return httpErr
+	err := p.listener.Close()
+	if err != nil {
+		p.logger.Error("Failed to close listener", "error", err)
+		return err
 	}
+
 	return nil
 }
 
+func (p *Server) isStarted() bool {
+	return p.started.Load()
+}
+
+func (p *Server) isStopped() bool {
+	return !p.started.Load()
+}
+
 // handleHTTP handles regular HTTP requests and CONNECT tunneling
 func (p *Server) handleHTTP(w http.ResponseWriter, r *http.Request) {
 	p.logger.Debug("handleHTTP called", "method", r.Method, "url", r.URL.String(), "host", r.Host)
@@ -479,13 +486,6 @@ func (p *Server) handleConnectionWithTLSDetection(conn net.Conn) {
 	}
 }
 
-// handleHTTPWithTLSTermination is the main handler (currently just delegates to regular HTTP)
-func (p *Server) handleHTTPWithTLSTermination(w http.ResponseWriter, r *http.Request) {
-	// This handler is not used when we do custom connection handling
-	// All traffic goes through handleConnectionWithTLSDetection
-	p.handleHTTP(w, r)
-}
-
 // connectionWrapper lets us "unread" the peeked byte
 type connectionWrapper struct {
 	net.Conn