From d171fca96df49a76f282879331993605b5a7a047 Mon Sep 17 00:00:00 2001
From: Andy Zhao <andyzhao@google.com>
Date: Wed, 22 Feb 2023 11:56:22 -0800
Subject: [PATCH] google: Add support for OAuth2 token exchange over mTLS

---
 google/default.go |  7 +++++++
 google/google.go  | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/google/default.go b/google/default.go
index 7ed02cd..bc6611d 100644
--- a/google/default.go
+++ b/google/default.go
@@ -6,6 +6,7 @@ package google
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -62,6 +63,12 @@ type CredentialsParams struct {
 
 	// PKCE is used to support PKCE flow. Optional for 3LO flow.
 	PKCE *authhandler.PKCEParams
+
+	// The OAuth2 TokenURL to use, which depends on whether mTLS is enabled. Optional.
+	TokenURL string
+
+	// The TLSConfig used for constructing an mTLS-enabled HTTP client. Optional.
+	TLSConfig *tls.Config
 }
 
 func (params CredentialsParams) deepCopy() CredentialsParams {
diff --git a/google/google.go b/google/google.go
index 8df0c49..1537a32 100644
--- a/google/google.go
+++ b/google/google.go
@@ -6,9 +6,12 @@ package google
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net"
+	"net/http"
 	"net/url"
 	"strings"
 	"time"
@@ -16,6 +19,7 @@ import (
 	"cloud.google.com/go/compute/metadata"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google/internal/externalaccount"
+	"golang.org/x/oauth2/internal"
 	"golang.org/x/oauth2/jwt"
 )
 
@@ -26,6 +30,9 @@ var Endpoint = oauth2.Endpoint{
 	AuthStyle: oauth2.AuthStyleInParams,
 }
 
+// MTLSTokenURL is Google's OAuth 2.0 default mTLS endpoint.
+const MTLSTokenURL = "https://oauth2.mtls.googleapis.com/token"
+
 // JWTTokenURL is Google's OAuth 2.0 token URL to use with the JWT flow.
 const JWTTokenURL = "https://oauth2.googleapis.com/token"
 
@@ -172,7 +179,14 @@ func (f *credentialsFile) tokenSource(ctx context.Context, params CredentialsPar
 			cfg.Endpoint.AuthURL = Endpoint.AuthURL
 		}
 		if cfg.Endpoint.TokenURL == "" {
-			cfg.Endpoint.TokenURL = Endpoint.TokenURL
+			if params.TokenURL != "" {
+				cfg.Endpoint.TokenURL = params.TokenURL
+			} else {
+				cfg.Endpoint.TokenURL = Endpoint.TokenURL
+			}
+		}
+		if params.TLSConfig != nil {
+			ctx = context.WithValue(ctx, internal.HTTPClient, customHTTPClient(params.TLSConfig))
 		}
 		tok := &oauth2.Token{RefreshToken: f.RefreshToken}
 		return cfg.TokenSource(ctx, tok), nil
@@ -275,3 +289,26 @@ func (cs computeSource) Token() (*oauth2.Token, error) {
 		"oauth2.google.serviceAccount": acct,
 	}), nil
 }
+
+// customHTTPClient constructs an HTTPClient using the provided tlsConfig, to support mTLS.
+func customHTTPClient(tlsConfig *tls.Config) *http.Client {
+	trans := baseTransport()
+	trans.TLSClientConfig = tlsConfig
+	return &http.Client{Transport: trans}
+}
+
+func baseTransport() *http.Transport {
+	return &http.Transport{
+		Proxy: http.ProxyFromEnvironment,
+		DialContext: (&net.Dialer{
+			Timeout:   30 * time.Second,
+			KeepAlive: 30 * time.Second,
+			DualStack: true,
+		}).DialContext,
+		MaxIdleConns:          100,
+		MaxIdleConnsPerHost:   100,
+		IdleConnTimeout:       90 * time.Second,
+		TLSHandshakeTimeout:   10 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+	}
+}