From 685ec81a16c29f3ae9e487042f954ed59adb97df Mon Sep 17 00:00:00 2001
From: Mohan Li <67390330+mohanli-ml@users.noreply.github.com>
Date: Wed, 19 Apr 2023 13:45:34 -0700
Subject: [PATCH] feat: add an option to enable DirectPath xDS (#1942)

Currently DirectPath xDS is enabled by using env. We want to also provide an option for this.

Java PR: googleapis/gax-java#1968
---
 internal/settings.go                    |  1 +
 option/internaloption/internaloption.go | 15 +++++++++++++++
 transport/grpc/dial.go                  | 15 ++++++++++++++-
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/internal/settings.go b/internal/settings.go
index 76efdb22772..3a3874df112 100644
--- a/internal/settings.go
+++ b/internal/settings.go
@@ -46,6 +46,7 @@ type DialSettings struct {
 	SkipValidation                bool
 	ImpersonationConfig           *impersonate.Config
 	EnableDirectPath              bool
+	EnableDirectPathXds           bool
 	AllowNonDefaultServiceAccount bool
 
 	// Google API system parameters. For more information please read:
diff --git a/option/internaloption/internaloption.go b/option/internaloption/internaloption.go
index cc7ebfe277b..3b8461d1da9 100644
--- a/option/internaloption/internaloption.go
+++ b/option/internaloption/internaloption.go
@@ -67,6 +67,21 @@ func (e enableDirectPath) Apply(o *internal.DialSettings) {
 	o.EnableDirectPath = bool(e)
 }
 
+// EnableDirectPathXds returns a ClientOption that overrides the default
+// DirectPath type. It is only valid when DirectPath is enabled.
+//
+// It should only be used internally by generated clients.
+// This is an EXPERIMENTAL API and may be changed or removed in the future.
+func EnableDirectPathXds() option.ClientOption {
+	return enableDirectPathXds(true)
+}
+
+type enableDirectPathXds bool
+
+func (x enableDirectPathXds) Apply(o *internal.DialSettings) {
+	o.EnableDirectPathXds = bool(x)
+}
+
 // AllowNonDefaultServiceAccount returns a ClientOption that overrides the default
 // requirement for using the default service account for DirectPath.
 //
diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go
index 9a9f57cc4e4..e1403e08ee6 100644
--- a/transport/grpc/dial.go
+++ b/transport/grpc/dial.go
@@ -164,7 +164,7 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C
 				grpcOpts = append(grpcOpts, timeoutDialerOption)
 			}
 			// Check if google-c2p resolver is enabled for DirectPath
-			if strings.EqualFold(os.Getenv(enableDirectPathXds), "true") {
+			if isDirectPathXdsUsed(o) {
 				// google-c2p resolver target must not have a port number
 				if addr, _, err := net.SplitHostPort(endpoint); err == nil {
 					endpoint = "google-c2p:///" + addr
@@ -251,6 +251,19 @@ func isDirectPathEnabled(endpoint string, o *internal.DialSettings) bool {
 	return true
 }
 
+func isDirectPathXdsUsed(o *internal.DialSettings) bool {
+	// Method 1: Enable DirectPath xDS by env;
+	if strings.EqualFold(os.Getenv(enableDirectPathXds), "true") {
+		return true
+	}
+	// Method 2: Enable DirectPath xDS by option;
+	if o.EnableDirectPathXds {
+		return true
+	}
+	return false
+
+}
+
 func isTokenSourceDirectPathCompatible(ts oauth2.TokenSource, o *internal.DialSettings) bool {
 	if ts == nil {
 		return false