From ceaa866219bd09735d89c4e66e48b90534c2ef8f Mon Sep 17 00:00:00 2001
From: Andy Zhao <andyzhao@google.com>
Date: Tue, 5 May 2020 16:55:31 -0700
Subject: [PATCH] google: Add support for 3-legged-OAuth using OAuth Client ID

Add OAuthClientTokenSource in google/google.go
Add DefaultAuthorizationHandler in authhandler.go
---
 google/authhandler.go | 23 +++++++++++++++++++++++
 google/google.go      | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 google/authhandler.go

diff --git a/google/authhandler.go b/google/authhandler.go
new file mode 100644
index 0000000..c92bac8
--- /dev/null
+++ b/google/authhandler.go
@@ -0,0 +1,23 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package google
+
+import (
+	"fmt"
+)
+
+const DefaultState = "state"
+
+// DefaultAuthorizationHandler is a commandline-based auth handler
+// that prints the auth URL on the console and prompts the user to
+// authorize in the browser and paste the auth code back via stdin.
+// When using this auth handler, DefaultState must be used.
+func DefaultAuthorizationHandler(authCodeUrl string) (string, string, error) {
+	fmt.Printf("Go to the following link in your browser:\n\n   %s\n\n", authCodeUrl)
+	fmt.Println("Enter verification code: ")
+	var code string
+	fmt.Scanln(&code)
+	return code, DefaultState, nil
+}
diff --git a/google/google.go b/google/google.go
index 81de32b..ebcabe5 100644
--- a/google/google.go
+++ b/google/google.go
@@ -207,3 +207,38 @@ func (cs computeSource) Token() (*oauth2.Token, error) {
 		"oauth2.google.serviceAccount": acct,
 	}), nil
 }
+
+// AuthorizationHandler is a 3-legged-OAuth helper that
+// prompts the user for OAuth consent at the specified Auth URL
+// and returns an auth code and state upon approval.
+type AuthorizationHandler func(string) (string, string, error)
+
+// OAuthClientTokenSource returns an oauth2.TokenSource that fetches access tokens
+// using 3-legged-OAuth workflow.
+// The provided oauth2.Config should be a full configuration containing AuthURL,
+// TokenURL, and scope.
+// An environment-specific AuthorizationHandler is used to obtain user consent.
+// Per OAuth protocol, a unique "state" string should be sent and verified
+// before token exchange to prevent CSRF attacks.
+func OAuthClientTokenSource(config oauth2.Config, ctx context.Context, authHandler AuthorizationHandler, state string) oauth2.TokenSource {
+	return oauth2.ReuseTokenSource(nil, oauthClientSource{config: config, ctx: ctx, authHandler: authHandler, state: state})
+}
+
+type oauthClientSource struct {
+	config      oauth2.Config
+	ctx         context.Context
+	authHandler AuthorizationHandler
+	state       string
+}
+
+func (ocs oauthClientSource) Token() (*oauth2.Token, error) {
+	url := ocs.config.AuthCodeURL(ocs.state)
+	code, state, err := ocs.authHandler(url)
+	if err != nil {
+		return nil, err
+	}
+	if state == ocs.state {
+		return ocs.config.Exchange(ocs.ctx, code)
+	}
+	return nil, errors.New("State mismatch in OAuth workflow.")
+}