From 39528763ae441fa965c8f0b1a5850cfd4b6d6fb0 Mon Sep 17 00:00:00 2001
From: Fuyang Liu <liufuyang@users.noreply.github.com>
Date: Mon, 12 Sep 2022 10:19:06 +0200
Subject: [PATCH] feat(cloudclient): support id_token with
 impersonated_service_account type

---
 cloudclient/dial.go | 74 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 73 insertions(+), 1 deletion(-)

diff --git a/cloudclient/dial.go b/cloudclient/dial.go
index fd37afa..5686182 100644
--- a/cloudclient/dial.go
+++ b/cloudclient/dial.go
@@ -3,11 +3,16 @@ package cloudclient
 import (
 	"context"
 	"crypto/x509"
+	"encoding/json"
 	"fmt"
 	"strconv"
 	"strings"
 	"time"
 
+	"go.einride.tech/cloudrunner/cloudzap"
+	"go.uber.org/zap"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
 	"google.golang.org/api/idtoken"
 	"google.golang.org/api/option"
 	"google.golang.org/grpc"
@@ -21,8 +26,52 @@ func DialService(ctx context.Context, target string, opts ...grpc.DialOption) (*
 	audience := "https://" + trimPort(target)
 	idTokenSource, err := idtoken.NewTokenSource(ctx, audience, option.WithAudiences(audience))
 	if err != nil {
-		return nil, fmt.Errorf("dial %s: %w", target, err)
+		logger, ok := cloudzap.GetLogger(ctx)
+		if !ok {
+			logger = zap.NewNop()
+		}
+		// Google's idtoken package does not support credential type other than `service_account`.
+		// This blocks local development with using `impersonated_service_account` type credentials. If that happens,
+		// we work it around by using our Application Default Credentials (which is impersonated already) to fetch
+		// an id_token on the fly.
+		// This however still blocks `authorized_user` type of credentials passing through.
+		// Related issue page: https://github.com/googleapis/google-api-go-client/issues/873
+		defaultCred, defaultCredErr := google.FindDefaultCredentials(ctx)
+		if defaultCredErr == nil {
+			var credTypeHolder struct {
+				Type string `json:"type"`
+			}
+			if jsonErr := json.Unmarshal(defaultCred.JSON, &credTypeHolder); jsonErr != nil {
+				// Ignoring jsonErr if it happens.
+				return nil, err
+			}
+			if credTypeHolder.Type != "impersonated_service_account" {
+				// We only patching the case where type of "impersonated_service_account" is used
+				// if not return original error.
+				return nil, err
+			}
+		} else {
+			// Ignoring defaultCredErr if it happens.
+			return nil, err
+		}
+
+		// Here we know the err is due to "impersonated_service_account" type of credential is used
+		// Just verify again with err's error message to be sure
+		if !strings.Contains(err.Error(), "impersonated_service_account") {
+			return nil, err
+		}
+
+		logger.Warn(
+			"Using Application Default Credentials to fetch id_token again. This should never happen in prod env.",
+			zap.String("ignored_error", err.Error()),
+		)
+		gts, err := google.DefaultTokenSource(ctx)
+		if err != nil {
+			return nil, err
+		}
+		idTokenSource = oauth2.ReuseTokenSource(nil, &idTokenSourceWrapper{TokenSource: gts})
 	}
+
 	systemCertPool, err := x509.SystemCertPool()
 	if err != nil {
 		return nil, fmt.Errorf("dial %s: %w", target, err)
@@ -62,3 +111,26 @@ func withDefaultPort(target string, port int) string {
 	}
 	return target
 }
+
+// idTokenSourceWrapper is an oauth2.TokenSource wrapper used for getting id_token for local development using
+// `authorized_user` type credentials
+// It takes the id_token from TokenSource and passes that on as a bearer token.
+type idTokenSourceWrapper struct {
+	TokenSource oauth2.TokenSource
+}
+
+func (s *idTokenSourceWrapper) Token() (*oauth2.Token, error) {
+	token, err := s.TokenSource.Token()
+	if err != nil {
+		return nil, err
+	}
+	idToken, ok := token.Extra("id_token").(string)
+	if !ok {
+		return nil, fmt.Errorf("token did not contain an id_token")
+	}
+	return &oauth2.Token{
+		AccessToken: idToken,
+		TokenType:   "Bearer",
+		Expiry:      token.Expiry,
+	}, nil
+}