From 0ca7dca97726252050774a4bff20d92ca5772331 Mon Sep 17 00:00:00 2001
From: yihuaz <yihuaz@google.com>
Date: Tue, 7 Sep 2021 09:12:01 -0700
Subject: [PATCH] oauth: Allow access to Google API regional endpoints via
 Google Default Credentials (#4713)

---
 credentials/oauth/oauth.go      | 19 ++++++++++-
 credentials/oauth/oauth_test.go | 60 +++++++++++++++++++++++++++++++++
 internal/credentials/util.go    |  4 ++-
 3 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 credentials/oauth/oauth_test.go

diff --git a/credentials/oauth/oauth.go b/credentials/oauth/oauth.go
index 852ae375cfc..c748fd21ce2 100644
--- a/credentials/oauth/oauth.go
+++ b/credentials/oauth/oauth.go
@@ -23,6 +23,7 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
+	"net/url"
 	"sync"
 
 	"golang.org/x/oauth2"
@@ -56,6 +57,16 @@ func (ts TokenSource) RequireTransportSecurity() bool {
 	return true
 }
 
+// removeServiceNameFromJWTURI removes RPC service name from URI.
+func removeServiceNameFromJWTURI(uri string) (string, error) {
+	parsed, err := url.Parse(uri)
+	if err != nil {
+		return "", err
+	}
+	parsed.Path = "/"
+	return parsed.String(), nil
+}
+
 type jwtAccess struct {
 	jsonKey []byte
 }
@@ -75,9 +86,15 @@ func NewJWTAccessFromKey(jsonKey []byte) (credentials.PerRPCCredentials, error)
 }
 
 func (j jwtAccess) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+	// Remove RPC service name from URI that will be used as audience
+	// in a self-signed JWT token. It follows https://google.aip.dev/auth/4111.
+	aud, err := removeServiceNameFromJWTURI(uri[0])
+	if err != nil {
+		return nil, err
+	}
 	// TODO: the returned TokenSource is reusable. Store it in a sync.Map, with
 	// uri as the key, to avoid recreating for every RPC.
-	ts, err := google.JWTAccessTokenSourceFromJSON(j.jsonKey, uri[0])
+	ts, err := google.JWTAccessTokenSourceFromJSON(j.jsonKey, aud)
 	if err != nil {
 		return nil, err
 	}
diff --git a/credentials/oauth/oauth_test.go b/credentials/oauth/oauth_test.go
new file mode 100644
index 00000000000..7e62ecb36c1
--- /dev/null
+++ b/credentials/oauth/oauth_test.go
@@ -0,0 +1,60 @@
+/*
+ *
+ * Copyright 2021 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package oauth
+
+import (
+	"strings"
+	"testing"
+)
+
+func checkErrorMsg(err error, msg string) bool {
+	if err == nil && msg == "" {
+		return true
+	} else if err != nil {
+		return strings.Contains(err.Error(), msg)
+	}
+	return false
+}
+
+func TestRemoveServiceNameFromJWTURI(t *testing.T) {
+	tests := []struct {
+		name         string
+		uri          string
+		wantedURI    string
+		wantedErrMsg string
+	}{
+		{
+			name:         "invalid URI",
+			uri:          "ht tp://foo.com",
+			wantedErrMsg: "first path segment in URL cannot contain colon",
+		},
+		{
+			name:      "valid URI",
+			uri:       "https://foo.com/go/",
+			wantedURI: "https://foo.com/",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got, err := removeServiceNameFromJWTURI(tt.uri); got != tt.wantedURI || !checkErrorMsg(err, tt.wantedErrMsg) {
+				t.Errorf("RemoveServiceNameFromJWTURI() = %s, %v, want %s, %v", got, err, tt.wantedURI, tt.wantedErrMsg)
+			}
+		})
+	}
+}
diff --git a/internal/credentials/util.go b/internal/credentials/util.go
index 55664fa46b8..f792fd22caf 100644
--- a/internal/credentials/util.go
+++ b/internal/credentials/util.go
@@ -18,7 +18,9 @@
 
 package credentials
 
-import "crypto/tls"
+import (
+	"crypto/tls"
+)
 
 const alpnProtoStrH2 = "h2"