dial.go

package cloudclient

import (
	"context"
	"crypto/x509"
	"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"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/keepalive"
)

// DialService dials another Cloud Run gRPC service with the default service account's RPC credentials.
func DialService(ctx context.Context, target string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	audience := "https://" + trimPort(target)
	idTokenSource, err := idtoken.NewTokenSource(ctx, audience, option.WithAudiences(audience))
	if err != nil {
		// Google's idtoken package does not support credential type other than `service_account`.
		// This blocks local development with using `authorized_user` type credentials. If that happens,
		// we work it around by using our Application Default Credentials to fetch an id_token on the fly.
		// Related issue page: https://github.com/googleapis/google-api-go-client/issues/873
		if !strings.HasPrefix(err.Error(), `idtoken: credential must be service_account, found`) {
			return nil, err
		}
		logger, ok := cloudzap.GetLogger(ctx)
		if !ok {
			logger = zap.NewNop()
		}
		logger.Sugar().Warnf("%s. Using Application Default Credentials to fetch id_token again. "+
			"This should never happen in prod env.", 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)
	}
	defaultOpts := []grpc.DialOption{
		grpc.WithPerRPCCredentials(
			&oauth.TokenSource{TokenSource: idTokenSource},
		),
		grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(systemCertPool, "")),
		// Enable connection keepalive to mitigate "connection reset by peer".
		// https://cloud.google.com/run/docs/troubleshooting
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                30 * time.Second,
			Timeout:             10 * time.Second,
			PermitWithoutStream: true,
		}),
	}
	conn, err := grpc.DialContext(ctx, withDefaultPort(target, 443), append(defaultOpts, opts...)...)
	if err != nil {
		return nil, fmt.Errorf("dial %s: %w", target, err)
	}
	return conn, nil
}

func trimPort(target string) string {
	parts := strings.Split(target, ":")
	if len(parts) == 1 {
		return target
	}
	return strings.Join(parts[:len(parts)-1], ":")
}

func withDefaultPort(target string, port int) string {
	parts := strings.Split(target, ":")
	if len(parts) == 1 {
		return target + ":" + strconv.Itoa(port)
	}
	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
}