From caea95689f82049822552cd649765335123831e0 Mon Sep 17 00:00:00 2001
From: Kui Xu <xmenxk@google.com>
Date: Mon, 14 Aug 2023 13:36:52 -0700
Subject: [PATCH] feat: ok to use S2A with override endpoint (#2114)

---
 internal/cba.go      |  4 ++--
 internal/cba_test.go | 11 +++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/internal/cba.go b/internal/cba.go
index 6923d3a716e..829383f55b5 100644
--- a/internal/cba.go
+++ b/internal/cba.go
@@ -274,8 +274,8 @@ func shouldUseS2A(clientCertSource cert.Source, settings *DialSettings) bool {
 	if !isGoogleS2AEnabled() {
 		return false
 	}
-	// If DefaultMTLSEndpoint is not set, skip S2A.
-	if settings.DefaultMTLSEndpoint == "" {
+	// If DefaultMTLSEndpoint is not set and no endpoint override, skip S2A.
+	if settings.DefaultMTLSEndpoint == "" && settings.Endpoint == "" {
 		return false
 	}
 	// If MTLS is not enabled for this endpoint, skip S2A.
diff --git a/internal/cba_test.go b/internal/cba_test.go
index 761d8e7d07e..d6a783740e5 100644
--- a/internal/cba_test.go
+++ b/internal/cba_test.go
@@ -290,6 +290,17 @@ func TestGetHTTPTransportConfigAndEndpoint(t *testing.T) {
 			testRegularEndpoint,
 			true,
 		},
+		{
+			"no client cert, S2A address not empty, override endpoint is set",
+			&DialSettings{
+				DefaultMTLSEndpoint: "",
+				Endpoint:            testOverrideEndpoint,
+			},
+			validConfigResp,
+			func() bool { return true },
+			testOverrideEndpoint,
+			false,
+		},
 		{
 			"no client cert, endpoint is MTLS enabled, S2A address not empty, custom HTTP client",
 			&DialSettings{