Tidy up, improve logging and add missing error handling.

Author: tecarter94

domain-proxy/common.go

@@ -8,7 +8,6 @@ import (
 	"os"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 )
 
@@ -20,7 +19,7 @@ const (
 	ConnectionTimeoutKey     = "CONNECTION_TIMEOUT"
 	DefaultConnectionTimeout = 1000 * time.Millisecond
 	IdleTimeoutKey           = "IDLE_TIMEOUT"
-	DefaultIdleTimeout       = 30000 * time.Millisecond
+	DefaultIdleTimeout       = 10000 * time.Millisecond
 )
 
 var Logger *log.Logger
@@ -29,53 +28,78 @@ func InitLogger(appName string) {
 	Logger = log.New(os.Stdout, appName+" ", log.LstdFlags|log.Lshortfile)
 }
 
-func BiDirectionalTransfer(leftConn, rightConn net.Conn, byteBufferSize int, idleTimeout time.Duration, executor *sync.WaitGroup) {
-	defer executor.Done()
-	defer CloseConnections(leftConn, rightConn)
+func BiDirectionalTransfer(leftConnection, rightConnection net.Conn, byteBufferSize int, idleTimeout time.Duration, connectionType string, connectionNo uint64) {
+	defer CloseConnection(leftConnection, rightConnection, connectionType, connectionNo)
 	done := make(chan struct{}, 2)
-	leftConn.SetDeadline(time.Now().Add(idleTimeout))
-	rightConn.SetDeadline(time.Now().Add(idleTimeout))
-	go Transfer(leftConn, rightConn, done, byteBufferSize, idleTimeout)
-	go Transfer(rightConn, leftConn, done, byteBufferSize, idleTimeout)
+	if err := leftConnection.SetDeadline(time.Now().Add(idleTimeout)); err != nil {
+		handleSetDeadlineError(leftConnection, err)
+		return
+	}
+	if err := rightConnection.SetDeadline(time.Now().Add(idleTimeout)); err != nil {
+		handleSetDeadlineError(rightConnection, err)
+		return
+	}
+	go Transfer(leftConnection, rightConnection, done, byteBufferSize, idleTimeout, connectionType, connectionNo)
+	go Transfer(rightConnection, leftConnection, done, byteBufferSize, idleTimeout, connectionType, connectionNo)
 	<-done
 	<-done
 }
 
-func Transfer(sourceConn, targetConn net.Conn, done chan struct{}, bufferSize int, idleTimeout time.Duration) {
+func Transfer(sourceConnection, targetConnection net.Conn, done chan struct{}, bufferSize int, idleTimeout time.Duration, connectionType string, connectionNo uint64) {
 	defer func() {
 		done <- struct{}{}
 	}()
 	buf := make([]byte, bufferSize)
 	for {
-		n, err := io.CopyBuffer(sourceConn, targetConn, buf)
-		if err != nil {
-			handleConnectionError(err)
+		if n, err := io.CopyBuffer(sourceConnection, targetConnection, buf); err != nil {
+			handleConnectionError(err, connectionType, connectionNo)
 			return
 		} else if n > 0 {
-			sourceConn.SetReadDeadline(time.Now().Add(idleTimeout))
-			targetConn.SetWriteDeadline(time.Now().Add(idleTimeout))
-			Logger.Printf("%d bytes transferred", n)
+			if err = sourceConnection.SetReadDeadline(time.Now().Add(idleTimeout)); err != nil {
+				handleSetDeadlineError(sourceConnection, err)
+				return
+			}
+			if err = targetConnection.SetWriteDeadline(time.Now().Add(idleTimeout)); err != nil {
+				handleSetDeadlineError(targetConnection, err)
+				return
+			}
+			Logger.Printf("%d bytes transferred for %s connection %d", n, connectionType, connectionNo)
 		}
 	}
 }
 
-func handleConnectionError(err error) {
+func handleSetDeadlineError(connection net.Conn, err error) {
+	Logger.Printf("Failed to set deadline: %v", err)
+	if err = connection.Close(); err != nil {
+		HandleConnectionCloseError(err)
+	}
+}
+
+func HandleConnectionCloseError(err error) {
+	Logger.Printf("Failed to close connection: %v", err)
+}
+
+func HandleListenerCloseError(err error) {
+	Logger.Printf("Failed to close listener: %v", err)
+}
+
+func handleConnectionError(err error, connectionType string, connectionNo uint64) {
 	var netErr net.Error
 	if errors.As(err, &netErr) && netErr.Timeout() {
-		Logger.Printf("Connection timed out")
+		Logger.Printf("%s connection %d timed out", connectionType, connectionNo)
 	} else if err != io.EOF {
-		Logger.Printf("Error using connection: %v", err)
+		Logger.Printf("Failed to transfer data using %s connection %d: %v", connectionType, connectionNo, err)
 	}
 }
 
-func CloseConnections(leftConn, rightConn net.Conn) {
-	if leftConn != nil {
-		leftConn.Close()
+func CloseConnection(leftConnection, rightConnection net.Conn, connectionType string, connectionNo uint64) {
+	if err := leftConnection.Close(); err != nil {
+		HandleConnectionCloseError(err)
 	}
-	if rightConn != nil {
-		rightConn.Close()
+	if err := rightConnection.Close(); err != nil {
+		HandleConnectionCloseError(err)
 	}
-	Logger.Println("Connections closed")
+	Logger.Printf("%s connection %d closed", connectionType, connectionNo)
 }
 
 func GetEnvVariable(key, defaultValue string) string {

domain-proxy/domain_proxy_client.go

@@ -5,7 +5,7 @@ import (
 	"net"
 	"os"
 	"os/signal"
-	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 )
@@ -14,17 +14,18 @@ const (
 	Localhost             = "localhost"
 	ServerHttpPortKey     = "SERVER_HTTP_PORT"
 	DefaultServerHttpPort = 8080
+	HttpToDomainSocket    = "HTTP <-> Domain Socket"
 )
 
 type DomainProxyClient struct {
-	domainSocket      string
-	serverHttpPort    int
-	byteBufferSize    int
-	connectionTimeout time.Duration
-	idleTimeout       time.Duration
-	listener          net.Listener
-	executor          *sync.WaitGroup
-	shutdownChan      chan struct{}
+	domainSocket          string
+	serverHttpPort        int
+	byteBufferSize        int
+	connectionTimeout     time.Duration
+	idleTimeout           time.Duration
+	httpConnectionCounter atomic.Uint64
+	listener              net.Listener
+	shutdownChan          chan struct{}
 }
 
 func NewDomainProxyClient(domainSocket string, serverHttpPort, byteBufferSize int, connectionTimeout, idleTimeout time.Duration) *DomainProxyClient {
@@ -34,7 +35,6 @@ func NewDomainProxyClient(domainSocket string, serverHttpPort, byteBufferSize in
 		byteBufferSize:    byteBufferSize,
 		connectionTimeout: connectionTimeout,
 		idleTimeout:       idleTimeout,
-		executor:          &sync.WaitGroup{},
 		shutdownChan:      make(chan struct{}),
 	}
 }
@@ -47,42 +47,49 @@ func (dpc *DomainProxyClient) Start() {
 	if err != nil {
 		Logger.Fatalf("Failed to start HTTP server: %v", err)
 	}
-	dpc.executor.Add(1)
 	go dpc.startClient()
 }
 
 func (dpc *DomainProxyClient) startClient() {
-	defer dpc.executor.Done()
-	Logger.Println("HTTP server listening on port", dpc.serverHttpPort)
+	Logger.Printf("HTTP server listening on port %d", dpc.serverHttpPort)
 	for {
-		serverConn, err := dpc.listener.Accept()
-		if err != nil {
+		if serverConnection, err := dpc.listener.Accept(); err != nil {
 			select {
 			case <-dpc.shutdownChan:
 				return
 			default:
-				Logger.Printf("Failed to accept connection: %v", err)
-				continue
+				Logger.Printf("Failed to accept server connection: %v", err)
 			}
+		} else {
+			go dpc.handleConnectionRequest(serverConnection)
 		}
-		domainConn, err := net.DialTimeout("unix", dpc.domainSocket, dpc.connectionTimeout)
-		if err != nil {
-			Logger.Printf("Failed to connect to domain socket: %v", err)
-			serverConn.Close()
-			continue
+	}
+}
+
+func (dpc *DomainProxyClient) handleConnectionRequest(serverConnection net.Conn) {
+	connectionNo := dpc.httpConnectionCounter.Add(1)
+	Logger.Printf("Handling %s Connection %d", HttpToDomainSocket, connectionNo)
+	startTime := time.Now()
+	domainConnection, err := net.DialTimeout("unix", dpc.domainSocket, dpc.connectionTimeout)
+	if err != nil {
+		Logger.Printf("Failed to connect to domain socket: %v", err)
+		if err = serverConnection.Close(); err != nil {
+			HandleConnectionCloseError(err)
 		}
-		dpc.executor.Add(1)
-		go BiDirectionalTransfer(serverConn, domainConn, dpc.byteBufferSize, dpc.idleTimeout, dpc.executor)
+		return
 	}
+	go func() {
+		BiDirectionalTransfer(serverConnection, domainConnection, dpc.byteBufferSize, dpc.idleTimeout, HttpToDomainSocket, connectionNo)
+		Logger.Printf("%s Connection %d ended after %d ms", HttpToDomainSocket, connectionNo, time.Since(startTime).Milliseconds())
+	}()
 }
 
 func (dpc *DomainProxyClient) Stop() {
 	Logger.Println("Shutting down domain proxy client...")
 	close(dpc.shutdownChan)
 	if err := dpc.listener.Close(); err != nil {
-		Logger.Printf("Error closing listener: %v", err)
+		HandleListenerCloseError(err)
 	}
-	//dpc.executor.Wait()
 }
 
 func GetServerHttpPort() int {
@@ -93,8 +100,8 @@ func main() {
 	InitLogger("Domain Proxy Client")
 	client := NewDomainProxyClient(GetDomainSocket(), GetServerHttpPort(), GetByteBufferSize(), GetConnectionTimeout(), GetIdleTimeout())
 	client.Start()
-	sigs := make(chan os.Signal, 1)
-	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
-	<-sigs
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+	<-signals
 	client.Stop()
 }