From f9a44934e9e12d01065e8cb9846c6991f7ab47d2 Mon Sep 17 00:00:00 2001
From: Patrick Jones <ithuriel@google.com>
Date: Wed, 10 Feb 2021 17:39:58 -0800
Subject: [PATCH] Added in fixes found during manual testing of Azure import.

Change-Id: Icf21d58732fdf4e3caaca015bc10d84613d0f423
---
 google/internal/externalaccount/basecredentials.go | 2 +-
 google/internal/externalaccount/impersonate.go     | 2 +-
 google/internal/externalaccount/sts_exchange.go    | 8 ++++++--
 google/internal/externalaccount/urlcredsource.go   | 3 +++
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/google/internal/externalaccount/basecredentials.go b/google/internal/externalaccount/basecredentials.go
index 57a5870..2eb5c8e 100644
--- a/google/internal/externalaccount/basecredentials.go
+++ b/google/internal/externalaccount/basecredentials.go
@@ -96,7 +96,7 @@ func (c *Config) parse(ctx context.Context) (baseCredentialSource, error) {
 	} else if c.CredentialSource.File != "" {
 		return fileCredentialSource{File: c.CredentialSource.File, Format: c.CredentialSource.Format}, nil
 	} else if c.CredentialSource.URL != "" {
-		return urlCredentialSource{URL: c.CredentialSource.URL, Format: c.CredentialSource.Format, ctx: ctx}, nil
+		return urlCredentialSource{URL: c.CredentialSource.URL, Headers: c.CredentialSource.Headers, Format: c.CredentialSource.Format, ctx: ctx}, nil
 	}
 	return nil, fmt.Errorf("oauth2/google: unable to parse credential source")
 }
diff --git a/google/internal/externalaccount/impersonate.go b/google/internal/externalaccount/impersonate.go
index 1d29c46..f848bff 100644
--- a/google/internal/externalaccount/impersonate.go
+++ b/google/internal/externalaccount/impersonate.go
@@ -64,7 +64,7 @@ func (its impersonateTokenSource) Token() (*oauth2.Token, error) {
 		return nil, fmt.Errorf("oauth2/google: unable to read body: %v", err)
 	}
 	if c := resp.StatusCode; c < 200 || c > 299 {
-		return nil, fmt.Errorf("oauth2/google: status code %d: %s", c, body)
+		return nil, fmt.Errorf("oauth2/google: status code %d: %s", c, string(body))
 	}
 
 	var accessTokenResp impersonateTokenResponse
diff --git a/google/internal/externalaccount/sts_exchange.go b/google/internal/externalaccount/sts_exchange.go
index 1a1c9b4..9ed4779 100644
--- a/google/internal/externalaccount/sts_exchange.go
+++ b/google/internal/externalaccount/sts_exchange.go
@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -63,9 +64,12 @@ func ExchangeToken(ctx context.Context, endpoint string, request *STSTokenExchan
 	}
 	defer resp.Body.Close()
 
-	bodyJson := json.NewDecoder(io.LimitReader(resp.Body, 1<<20))
+	body, err := ioutil.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	if c := resp.StatusCode; c < 200 || c > 299 {
+		return nil, fmt.Errorf("oauth2/google: status code %d: %s", c, string(body))
+	}
 	var stsResp STSTokenExchangeResponse
-	err = bodyJson.Decode(&stsResp)
+	json.Unmarshal(body, &stsResp)
 	if err != nil {
 		return nil, fmt.Errorf("oauth2/google: failed to unmarshal response body from Secure Token Server: %v", err)
 
diff --git a/google/internal/externalaccount/urlcredsource.go b/google/internal/externalaccount/urlcredsource.go
index b0d5d35..dce704e 100644
--- a/google/internal/externalaccount/urlcredsource.go
+++ b/google/internal/externalaccount/urlcredsource.go
@@ -43,6 +43,9 @@ func (cs urlCredentialSource) subjectToken() (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("oauth2/google: invalid body in subject token URL query: %v", err)
 	}
+	if c := resp.StatusCode; c < 200 || c > 299 {
+		return "", fmt.Errorf("oauth2/google: status code %d: %s", c, string(tokenBytes))
+	}
 
 	switch cs.Format.Type {
 	case "json":