From 8f90c1b16d1b6e288f373befc45c4a82aa526948 Mon Sep 17 00:00:00 2001
From: "Tanishq Dubey (aider)" <dubey@dws.rip>
Date: Sun, 18 May 2025 10:51:06 -0400
Subject: [PATCH] fix: update TLS configuration to use leader hostname and
 custom dialer

---
 internal/agent/agent.go | 49 ++++++++++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index a62f0c4..17802ca 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"runtime"
@@ -115,28 +116,40 @@ func (a *Agent) SetupMTLSClient() error {
 		Certificates: []tls.Certificate{cert},
 		RootCAs:      caCertPool,
 		MinVersion:   tls.VersionTLS12,
-		// Skip hostname verification since we're using IP addresses
-		// and the leader cert is issued for leader.kat.cluster.local
-		InsecureSkipVerify: true,
-		// Custom verification to still validate the certificate chain
-		// but ignore the hostname mismatch
-		VerifyPeerCertificate: func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
-			// Skip verification if there are no chains (shouldn't happen with our config)
-			if len(verifiedChains) == 0 {
-				return fmt.Errorf("no verified chains provided")
-			}
-			
-			// The certificate chain was already verified against our CA by the TLS stack
-			// We just need to check that the leaf cert was issued by our trusted CA
-			// which is already done by the time this callback is called
-			return nil
-		},
 	}
 
 	// Create HTTP client with TLS configuration
 	a.client = &http.Client{
 		Transport: &http.Transport{
 			TLSClientConfig: tlsConfig,
+			// Override the dial function to map any hostname to the leader's IP
+			DialTLS: func(network, addr string) (net.Conn, error) {
+				// Extract host and port from addr
+				host, port, err := net.SplitHostPort(addr)
+				if err != nil {
+					return nil, err
+				}
+				
+				// Extract host and port from LeaderAPI
+				leaderHost, leaderPort, err := net.SplitHostPort(a.LeaderAPI)
+				if err != nil {
+					return nil, err
+				}
+				
+				// Use the leader's IP but keep the original port
+				dialAddr := net.JoinHostPort(leaderHost, port)
+				
+				// For logging purposes
+				log.Printf("Dialing %s instead of %s", dialAddr, addr)
+				
+				// Create the TLS connection
+				conn, err := tls.Dial(network, dialAddr, tlsConfig)
+				if err != nil {
+					return nil, err
+				}
+				
+				return conn, nil
+			},
 		},
 		Timeout: 10 * time.Second,
 	}
@@ -198,8 +211,8 @@ func (a *Agent) sendHeartbeat() error {
 		return fmt.Errorf("failed to marshal node status: %w", err)
 	}
 
-	// Construct URL
-	url := fmt.Sprintf("https://%s/v1alpha1/nodes/%s/status", a.LeaderAPI, a.NodeName)
+	// Construct URL - use leader.kat.cluster.local as hostname to match certificate
+	url := fmt.Sprintf("https://leader.kat.cluster.local/v1alpha1/nodes/%s/status", a.NodeName)
 
 	// Create request
 	req, err := http.NewRequest("POST", url, bytes.NewBuffer(statusJSON))