OAuthWeb Credential Flow Should Run a Local Server (#20)

The OAuthWebCredentialFlow should run a local server that the OAuth web
flow can redirect to to automatically pass the OAuth token.

This code is based a lot on the OIDCWebFlow.

gcp/credentials.go

@@ -1,14 +1,23 @@
-// package gcp provides utilities for working with GCP
+// Package gcp provides utilities for working with GCP
 package gcp
 
 import (
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
+
+	"github.com/gorilla/mux"
+	"github.com/jlewi/monogo/networking"
+	"github.com/jlewi/monogo/oauthutil"
+	"github.com/jlewi/p22h/backend/api"
+	"github.com/jlewi/p22h/backend/pkg/debug"
+	"github.com/pkg/browser"
 
 	"github.com/jlewi/monogo/files"
 
@@ -25,21 +34,35 @@ import (
 const (
 	// CredentialDirPermMode unix permission max suitable for directory storing credentials
 	CredentialDirPermMode = 0700
+	authCallbackUrl       = "/auth/callback"
+	authStartPrefix       = "/auth/start"
 )
 
 // CredentialHelper defines an interface for getting tokens.
 type CredentialHelper interface {
-	//GetTokenAndConfig() (*TokenAndConfig, error)
 	GetTokenSource(ctx context.Context) (oauth2.TokenSource, error)
 
 	// GetOAuthConfig returns the OAuth2 client configuration
 	GetOAuthConfig() *oauth2.Config
 }
 
-// WebFlowHelper helps get credentials using the webflow.
+// WebFlowHelper helps get credentials using the webflow. It is intended for desktop applications.
+// It runs a local server to handle the callback from the OAuth server to get the authorization code and return
+// a token source.
+//
+// References: https://developers.google.com/identity/protocols/oauth2/native-app#request-parameter-redirect_uri
+// GCP still supports using the loopback device 127.0.0.1 for OAuth credentials for desktop applications.
+// It looks like in that case you don't actually have to specify your redirect URI when configuring the OAuth Client
+// in the developer console. However, when you specify your OAuth configuration in the code you need to specify the
+// redirect URI and it needs to be 127.0.0.1 not localhost.
 type WebFlowHelper struct {
-	config *oauth2.Config
-	Log    logr.Logger
+	config   *oauth2.Config
+	Log      logr.Logger
+	host     string
+	handlers *oauthutil.OAuthHandlers
+	// Server to handle callback
+	srv *http.Server
+	c   chan tokenSourceOrError
 }
 
 // NewWebFlowHelper constructs a new web flow helper. oAuthClientFile should be the path to a credentials.json
@@ -79,35 +102,193 @@ func NewWebFlowHelper(oAuthClientFile string, scopes []string) (*WebFlowHelper,
 	if err != nil {
 		return nil, errors.Wrapf(err, "Unable to parse client secret file to config")
 	}
+
+	port, err := networking.GetFreePort()
+	if err != nil {
+		return nil, errors.Wrapf(err, "Failed to get free port")
+	}
+	// Host must match the same as the redirect URI i.e. we can't use localhost in one place and the loopback in
+	// another. This because when we open authStartPrefix in the browser it sets a cookie for the domain which will
+	// be host. We then read that cookie when we get redirected back to the callback URL. If the cookie is set
+	// on localhost then it won't be accessible when we get to 127.0.0.1 and the flow will fail.
+	host := fmt.Sprintf("127.0.0.1:%d", port)
+	// We need to use 127.0.0.1 because that's what GCP expects for desktop. If we use localhost it won't work.
+	// https://developers.google.com/identity/protocols/oauth2/native-app#request-parameter-redirect_uri
+	config.RedirectURL = fmt.Sprintf("http://127.0.0.1:%d%v", port, authCallbackUrl)
+	handlers, err := oauthutil.NewOAuthHandlers(*config)
+	if err != nil {
+		return nil, err
+	}
+
 	return &WebFlowHelper{
-		config: config,
-		Log:    zapr.NewLogger(zap.L()),
+		config:   config,
+		Log:      zapr.NewLogger(zap.L()),
+		handlers: handlers,
+		host:     host,
+		c:        make(chan tokenSourceOrError, 10),
 	}, nil
 }
 
 func (h *WebFlowHelper) GetOAuthConfig() *oauth2.Config {
 	return h.config
 }
 
+// Run runs the flow to create a tokensource.
+// It starts a server in order to provide a callback that the OAuthFlow can redirect to in order to pass the
+// authorization code.
+// The server is shutdown after the flow is complete. Since the flow should return a refresh token
+// it shouldn't be necessary to keep it running.
+func (h *WebFlowHelper) Run() (oauth2.TokenSource, error) {
+	log := zapr.NewLogger(zap.L())
+
+	go func() {
+		h.startAndBlock()
+	}()
+	log.Info("Waiting for OAuth server to be ready")
+	if err := h.waitForReady(); err != nil {
+		return nil, err
+	}
+	authURL := h.AuthStartURL()
+	log.Info("Opening URL to start Auth Flow", "URL", authURL)
+	if err := browser.OpenURL(authURL); err != nil {
+		log.Error(err, "Failed to open URL in browser; open it manually", "url", authURL)
+		// TODO(jeremy): How do we scan it it in? Should we fall back to calling GetTokenSource?
+		fmt.Printf("Go to the following link in your browser to complete  the OAuth flow: %v\n", authURL)
+	}
+	// Wait for the token source
+	log.Info("Waiting for OAuth flow to complete")
+
+	defer func() {
+		log.Info("Shutting OAuth server down")
+		err := h.srv.Shutdown(context.Background())
+		if err != nil {
+			log.Error(err, "There was a problem shutting the OAuth server down")
+		}
+	}()
+	select {
+	case tsOrError := <-h.c:
+		if tsOrError.err != nil {
+			return nil, errors.Wrapf(tsOrError.err, "OAuth flow didn't complete successfully")
+		}
+		log.Info("OAUth flow completed")
+		return tsOrError.ts, nil
+	case <-time.After(3 * time.Minute):
+		return nil, errors.New("Timeout waiting for OIDC flow to complete")
+	}
+}
+
+// startAndBlock starts the server and blocks.
+func (h *WebFlowHelper) startAndBlock() {
+	log := zapr.NewLogger(zap.L())
+
+	router := mux.NewRouter().StrictSlash(true)
+
+	router.HandleFunc(authStartPrefix, h.handleStartWebFlow)
+	router.HandleFunc("/healthz", h.HealthCheck)
+	router.HandleFunc(authCallbackUrl, h.handleAuthCallback)
+
+	router.NotFoundHandler = http.HandlerFunc(h.NotFoundHandler)
+
+	log.Info("OAuth server is running", "address", h.Address())
+
+	h.srv = &http.Server{Addr: h.host, Handler: router}
+
+	err := h.srv.ListenAndServe()
+
+	// ListenAndServe will return ErrServerClosed when the server is shutdown.
+	if err != nil && !errors.Is(err, http.ErrServerClosed) {
+		log.Error(err, "OAuth server returned error")
+	}
+	log.Info("OAuth server has been shutdown")
+}
+
+// waitForReady waits until the server is health.
+func (h *WebFlowHelper) waitForReady() error {
+	endTime := time.Now().Add(3 * time.Minute)
+	for time.Now().Before(endTime) {
+
+		r, err := http.Get(h.Address() + "/healthz")
+		if err == nil && r.StatusCode == http.StatusOK {
+			return nil
+		}
+		time.Sleep(5 * time.Second)
+	}
+	return errors.New("timeout waiting for server to be healthy")
+}
+
+// handleStartWebFlow kicks off the OAuthWebFlow.
+// It was copied from: https://github.com/coreos/go-oidc/blob/2cafe189143f4a454e8b4087ef892be64b1c77df/example/idtoken/app.go#L65
+// It sets some cookies before redirecting to the OAuth provider's URL for obtaining an authorization code.
+func (h *WebFlowHelper) handleStartWebFlow(w http.ResponseWriter, r *http.Request) {
+	log := zapr.NewLogger(zap.L())
+	// N.B. we currently ignore the cookie and state because we run thisflow in a CLI/application. The implicit
+	// assumption is that a single user is going through the flow a single time so we don't need to use
+	// cookie and state to keep track of the user's session. We also don't need to use the cookie
+	// to keep track of the page the user was visiting before they started the flow.
+	_, err := h.handlers.RedirectToAuthURL(w, r)
+	if err != nil {
+		log.Error(err, "Failed to handle auth start")
+	}
+}
+
+func (h *WebFlowHelper) handleAuthCallback(w http.ResponseWriter, r *http.Request) {
+	log := zapr.NewLogger(zap.L())
+	_, ts, err := h.handlers.HandleAuthCode(w, r)
+	if err != nil {
+		log.Error(err, "Failed to handle auth callback")
+		h.c <- tokenSourceOrError{err: err}
+		return
+	}
+
+	h.c <- tokenSourceOrError{ts: ts}
+	if _, err := w.Write([]byte("OAuth flow completed; you can close this window and return to your application")); err != nil {
+		log.Error(err, "Failed to write response")
+	}
+}
+
 // GetTokenSource requests a token from the web, then returns the retrieved token.
+// TODO(jeremy): Deprecate this method in favor of Run.
 func (h *WebFlowHelper) GetTokenSource(ctx context.Context) (oauth2.TokenSource, error) {
-	authURL := h.config.AuthCodeURL("state-token", oauth2.AccessTypeOffline)
+	return h.Run()
+}
 
-	// TODO(jlewi): How to open it automatically?
-	fmt.Printf("Go to the following link in your browser then type the "+
-		"authorization code: \n%v\n", authURL)
+// AuthStartURL returns the URL to kickoff the oauth login flow.
+func (s *WebFlowHelper) AuthStartURL() string {
+	return s.Address() + authStartPrefix
+}
 
-	var authCode string
-	if _, err := fmt.Scan(&authCode); err != nil {
-		return nil, errors.Wrapf(err, "Unable to read authorization code")
+func (h *WebFlowHelper) Address() string {
+	return fmt.Sprintf("http://%v", h.host)
+}
+
+func (h *WebFlowHelper) writeStatus(w http.ResponseWriter, message string, code int) {
+	log := zapr.NewLogger(zap.L())
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(code)
+
+	resp := api.RequestStatus{
+		Kind:    "RequestStatus",
+		Message: message,
+		Code:    code,
 	}
 
-	tok, err := h.config.Exchange(context.TODO(), authCode)
-	if err != nil {
-		return nil, errors.Wrapf(err, "Unable to retrieve token from web")
+	enc := json.NewEncoder(w)
+	if err := enc.Encode(resp); err != nil {
+		log.Error(err, "Failed to marshal RequestStatus", "RequestStatus", resp, "code", code)
 	}
 
-	return h.config.TokenSource(ctx, tok), nil
+	if code != http.StatusOK {
+		caller := debug.ThisCaller()
+		log.Info("HTTP error", "RequestStatus", resp, "code", code, "caller", caller)
+	}
+}
+
+func (h *WebFlowHelper) HealthCheck(w http.ResponseWriter, r *http.Request) {
+	h.writeStatus(w, "OAuth server is running", http.StatusOK)
+}
+
+func (h *WebFlowHelper) NotFoundHandler(w http.ResponseWriter, r *http.Request) {
+	h.writeStatus(w, fmt.Sprintf("OAuth server doesn't handle the path; url: %v", r.URL), http.StatusNotFound)
 }
 
 // TokenCache defines an interface for caching tokens
@@ -210,3 +391,8 @@ func (c *CachedCredentialHelper) GetTokenSource(ctx context.Context) (oauth2.Tok
 	ts := c.CredentialHelper.GetOAuthConfig().TokenSource(ctx, tok)
 	return ts, nil
 }
+
+type tokenSourceOrError struct {
+	ts  oauth2.TokenSource
+	err error
+}

gcp/credentials_test.go

@@ -0,0 +1,35 @@
+package gcp
+
+import (
+	"os"
+	"testing"
+
+	"go.uber.org/zap"
+	"google.golang.org/api/gmail/v1"
+)
+
+func Test_WebFlowHelepr(t *testing.T) {
+	if os.Getenv("GITHUB_ACTIONS") != "" {
+		t.Skip("Skipping test in GitHub Actions")
+	}
+	dLog, err := zap.NewDevelopmentConfig().Build()
+	if err != nil {
+		t.Fatalf("Error creating logger: %v", err)
+	}
+	zap.ReplaceGlobals(dLog)
+
+	clientSecret := "/Users/jlewi/secrets/gctl.oauthclientid.foyle-dev.json"
+	flow, err := NewWebFlowHelper(clientSecret, []string{gmail.GmailReadonlyScope})
+	if err != nil {
+		t.Fatalf("Error creating web flow helper: %v", err)
+	}
+
+	ts, err := flow.Run()
+	if err != nil {
+		t.Fatalf("Error getting token source: %v", err)
+	}
+
+	if _, err := ts.Token(); err != nil {
+		t.Fatalf("Error getting token: %v", err)
+	}
+}

oauthutil/oidc_webflow.go

@@ -160,13 +160,10 @@ func (s *OIDCWebFlowServer) Run() (*IDTokenSource, error) {
 			return nil, errors.Wrapf(tsOrError.err, "OIDC flow didn't complete successfully")
 		}
 		log.Info("OIDC flow completed")
-		// TODO(jeremy): This is a hack to deal with a race condition in which the server starts shutting down
-		// but oauth flow still sends it some request.
 		return tsOrError.ts, nil
 	case <-time.After(3 * time.Minute):
 		return nil, errors.New("Timeout waiting for OIDC flow to complete")
 	}
-
 }
 
 // startAndBlock starts the server and blocks.
@@ -187,7 +184,7 @@ func (s *OIDCWebFlowServer) startAndBlock() {
 
 	err := s.srv.ListenAndServe()
 
-	if err != nil {
+	if err != nil && !errors.Is(err, http.ErrServerClosed) {
 		log.Error(err, "OIDCWebFlowServer returned error")
 	}
 	log.Info("OIDC server has been shutdown")

oauthutil/server.go

@@ -101,7 +101,7 @@ func (s *OAuthHandlers) HandleAuthCode(w http.ResponseWriter, r *http.Request) (
 	}
 	actual := r.URL.Query().Get("state")
 	if actual != state.Value {
-		s.log.Info("state dind't match", "got", actual, "want", state.Value)
+		s.log.Info("state didn't match", "got", actual, "want", state.Value)
 		http.Error(w, "state did not match", http.StatusBadRequest)
 		return "", nil, err
 	}