feat: completed cli for gateway

cli/packages/gateway/connection.go

@@ -3,7 +3,7 @@ package gateway
 import (
 	"bufio"
 	"bytes"
-	"fmt"
+	"errors"
 	"io"
 	"net"
 	"sync"
@@ -20,6 +20,9 @@ func handleConnection(conn net.Conn) {
 	for {
 		msg, err := reader.ReadBytes('\n')
 		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return
+			}
 			log.Error().Msgf("Error reading command: %s", err)
 			return
 		}
@@ -30,9 +33,7 @@ func handleConnection(conn net.Conn) {
 		switch string(cmd) {
 		case "FORWARD-TCP":
 			proxyAddress := string(bytes.Split(args, []byte(" "))[0])
-			fmt.Println(proxyAddress)
 			destTarget, err := net.Dial("tcp", proxyAddress)
-			fmt.Println(err)
 			if err != nil {
 				log.Error().Msgf("Failed to connect to target: %v", err)
 				return
@@ -56,12 +57,13 @@ func handleConnection(conn net.Conn) {
 			}
 
 			CopyData(conn, destTarget)
-			break
+			return
 		case "PING":
-			conn.Write([]byte("PONG\n"))
+			conn.Write([]byte("PONG"))
+			return
 		default:
 			log.Error().Msgf("Unknown command: %s", string(cmd))
-			break
+			return
 		}
 	}
 }
@@ -71,39 +73,33 @@ type CloseWrite interface {
 }
 
 func CopyData(src, dst net.Conn) {
-	// Create a WaitGroup to wait for both copy operations
 	var wg sync.WaitGroup
 	wg.Add(2)
 
-	// Start copying in both directions
-	go func() {
+	copyAndClose := func(dst, src net.Conn, done chan<- bool) {
 		defer wg.Done()
-		if _, err := io.Copy(dst, src); err != nil {
-			log.Error().Msgf("Error copying postgres->client: %v", err)
+		_, err := io.Copy(dst, src)
+		if err != nil && !errors.Is(err, io.EOF) {
+			log.Error().Msgf("Copy error: %v", err)
 		}
 
-		if e, ok := dst.(CloseWrite); ok {
-			log.Print("Closing dst")
-			e.CloseWrite()
-		} else {
+		// Signal we're done writing
+		done <- true
 
-			log.Print("Not closed")
+		// Half close the connection if possible
+		if c, ok := dst.(CloseWrite); ok {
+			c.CloseWrite()
 		}
-	}()
+	}
 
-	go func() {
-		defer wg.Done()
-		if _, err := io.Copy(src, dst); err != nil {
-			log.Error().Msgf("Error copying client->postgres: %v", err)
-		}
-		if e, ok := src.(CloseWrite); ok {
-			log.Print("Closing src")
-			e.CloseWrite()
-		} else {
-			log.Print("Not closed")
-		}
-	}()
+	done1 := make(chan bool, 1)
+	done2 := make(chan bool, 1)
+
+	go copyAndClose(dst, src, done1)
+	go copyAndClose(src, dst, done2)
 
 	// Wait for both copies to complete
+	<-done1
+	<-done2
 	wg.Wait()
 }

cli/packages/gateway/gateway.go

@@ -48,16 +48,15 @@ func (g *Gateway) ConnectWithRelay() error {
 		return err
 	}
 
-	// Dial TURN Server
-	conn, err := net.Dial("tcp", relayDetails.TurnServerAddress)
+	turnServerAddr, err := net.ResolveTCPAddr("tcp", relayDetails.TurnServerAddress)
 	if err != nil {
-		return fmt.Errorf("Failed to connect with relay server: %w", err)
+		return fmt.Errorf("Failed to resolve TURN server address: %w", err)
 	}
 
-	if tcpConn, ok := conn.(*net.TCPConn); ok {
-		tcpConn.SetKeepAlive(true)
-		tcpConn.SetKeepAlivePeriod(10 * time.Second)
-		tcpConn.SetNoDelay(true)
+	// Dial TURN Server
+	conn, err := net.DialTCP("tcp", nil, turnServerAddr)
+	if err != nil {
+		return fmt.Errorf("Failed to connect with relay server: %w", err)
 	}
 
 	// Start a new TURN Client and wrap our net.Conn in a STUNConn
@@ -77,11 +76,6 @@ func (g *Gateway) ConnectWithRelay() error {
 		return fmt.Errorf("Failed to create relay client: %w", err)
 	}
 
-	err = client.Listen()
-	if err != nil {
-		return fmt.Errorf("Failed to listen to relay server: %w", err)
-	}
-
 	g.config = &GatewayConfig{
 		TurnServerUsername: relayDetails.TurnServerUsername,
 		TurnServerPassword: relayDetails.TurnServerPassword,
@@ -99,6 +93,11 @@ func (g *Gateway) ConnectWithRelay() error {
 
 func (g *Gateway) Listen() error {
 	defer g.client.Close()
+	err := g.client.Listen()
+	if err != nil {
+		return fmt.Errorf("Failed to listen to relay server: %w", err)
+	}
+	log.Info().Msg("Connected with relay")
 	// Allocate a relay socket on the TURN server. On success, it
 	// will return a net.PacketConn which represents the remote
 	// socket.
@@ -106,6 +105,7 @@ func (g *Gateway) Listen() error {
 	if err != nil {
 		return fmt.Errorf("Failed to allocate relay connection: %w", err)
 	}
+
 	defer func() {
 		if closeErr := relayNonTlsConn.Close(); closeErr != nil {
 			log.Error().Msgf("Failed to close connection: %s", closeErr)
@@ -148,7 +148,6 @@ func (g *Gateway) Listen() error {
 		return fmt.Errorf("failed to parse cert: %s", err)
 	}
 
-	fmt.Println(relayNonTlsConn.Addr().String())
 	caCertPool := x509.NewCertPool()
 	caCertPool.AppendCertsFromPEM([]byte(gatewayCert.CertificateChain))
 
@@ -159,8 +158,8 @@ func (g *Gateway) Listen() error {
 		ClientAuth:   tls.RequireAndVerifyClientCert,
 	})
 
+	log.Info().Msg("Connector started successfully")
 	for {
-		log.Info().Msg("Connector started successfully")
 		// Accept new relay connection
 		conn, err := relayConn.Accept()
 		if err != nil {
@@ -168,6 +167,31 @@ func (g *Gateway) Listen() error {
 			continue
 		}
 
+		tlsConn, ok := conn.(*tls.Conn)
+		if !ok {
+			log.Error().Msg("Failed to convert to TLS connection")
+			conn.Close()
+			continue
+		}
+
+		err = tlsConn.Handshake()
+		if err != nil {
+			log.Error().Msgf("TLS handshake failed: %v", err)
+			conn.Close()
+			continue
+		}
+
+		// Get connection state which contains certificate information
+		state := tlsConn.ConnectionState()
+		if len(state.PeerCertificates) > 0 {
+			organizationUnit := state.PeerCertificates[0].Subject.OrganizationalUnit
+			commonName := state.PeerCertificates[0].Subject.CommonName
+			if organizationUnit[0] != "gateway-client" && commonName != "cloud" {
+				log.Error().Msgf("Client certificate verification failed. Received %s, %s", organizationUnit, commonName)
+				continue
+			}
+		}
+
 		// Handle the connection in a goroutine
 		go handleConnection(conn)
 	}