From a4a909a20b0f868de4870294e200e803f61589f7 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 19 Feb 2024 08:56:58 -0800
Subject: [PATCH] prober: add TLS probe constructor to split dial addr from
 cert name

So we can probe load balancers by their unique DNS name but without
asking for that cert name.

Updates tailscale/corp#13050

Change-Id: Ie4c0a2f951328df64281ed1602b4e624e3c8cf2e
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 prober/tls.go      | 29 +++++++++++++++++++----------
 prober/tls_test.go |  2 +-
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/prober/tls.go b/prober/tls.go
index 97c53298e..db25c9fd5 100644
--- a/prober/tls.go
+++ b/prober/tls.go
@@ -12,6 +12,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/netip"
 	"time"
 
 	"github.com/pkg/errors"
@@ -23,25 +24,33 @@ const expiresSoon = 7 * 24 * time.Hour // 7 days from now
 
 // TLS returns a Probe that healthchecks a TLS endpoint.
 //
-// The ProbeFunc connects to a hostname (host:port string), does a TLS
+// The ProbeFunc connects to a hostPort (host:port string), does a TLS
 // handshake, verifies that the hostname matches the presented certificate,
 // checks certificate validity time and OCSP revocation status.
-func TLS(hostname string) ProbeFunc {
+func TLS(hostPort string) ProbeFunc {
 	return func(ctx context.Context) error {
-		return probeTLS(ctx, hostname)
+		certDomain, _, err := net.SplitHostPort(hostPort)
+		if err != nil {
+			return err
+		}
+		return probeTLS(ctx, certDomain, hostPort)
 	}
 }
 
-func probeTLS(ctx context.Context, hostname string) error {
-	host, _, err := net.SplitHostPort(hostname)
-	if err != nil {
-		return err
+// TLSWithIP is like TLS, but dials the provided dialAddr instead
+// of using DNS resolution. The certDomain is the expected name in
+// the cert (and the SNI name to send).
+func TLSWithIP(certDomain string, dialAddr netip.AddrPort) ProbeFunc {
+	return func(ctx context.Context) error {
+		return probeTLS(ctx, certDomain, dialAddr.String())
 	}
+}
 
-	dialer := &tls.Dialer{Config: &tls.Config{ServerName: host}}
-	conn, err := dialer.DialContext(ctx, "tcp", hostname)
+func probeTLS(ctx context.Context, certDomain string, dialHostPort string) error {
+	dialer := &tls.Dialer{Config: &tls.Config{ServerName: certDomain}}
+	conn, err := dialer.DialContext(ctx, "tcp", dialHostPort)
 	if err != nil {
-		return fmt.Errorf("connecting to %q: %w", hostname, err)
+		return fmt.Errorf("connecting to %q: %w", dialHostPort, err)
 	}
 	defer conn.Close()
 
diff --git a/prober/tls_test.go b/prober/tls_test.go
index 6c6601ee0..5bfb739db 100644
--- a/prober/tls_test.go
+++ b/prober/tls_test.go
@@ -85,7 +85,7 @@ func TestTLSConnection(t *testing.T) {
 	srv.StartTLS()
 	defer srv.Close()
 
-	err = probeTLS(context.Background(), srv.Listener.Addr().String())
+	err = probeTLS(context.Background(), "fail.example.com", srv.Listener.Addr().String())
 	// The specific error message here is platform-specific ("certificate is not trusted"
 	// on macOS and "certificate signed by unknown authority" on Linux), so only check
 	// that it contains the word 'certificate'.