From 08bbd1d1933ab638662914128b2f5ccd0efbd257 Mon Sep 17 00:00:00 2001
From: TakumiHaruta <tharuta919@gmail.com>
Date: Mon, 27 Feb 2023 09:05:08 +0900
Subject: [PATCH] feat: add static token option in client

---
 pkg/client/middleware/auth.go | 19 +++++++++++---
 pkg/client/zitadel/client.go  | 49 ++++++++++++++++++++++++++---------
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/pkg/client/middleware/auth.go b/pkg/client/middleware/auth.go
index 06616cc..f40c32e 100644
--- a/pkg/client/middleware/auth.go
+++ b/pkg/client/middleware/auth.go
@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"context"
+	"fmt"
 	"os"
 
 	"github.com/zitadel/oidc/pkg/client/profile"
@@ -42,10 +43,10 @@ func JWTProfileFromKeyAndUserID(key []byte, keyID, userID string) JWTProfileToke
 	}
 }
 
-//NewAuthenticator creates an interceptor which authenticates a service account with a provided JWT Profile (using a key.json either as file or data).
+//NewJWTProfileAuthenticator creates an interceptor which authenticates a service account with a provided JWT Profile (using a key.json either as file or data).
 //There returned token will be used for authorization in all calls
 //if expired, the token will be automatically refreshed
-func NewAuthenticator(issuer string, jwtProfileTokenSource JWTProfileTokenSource, scopes ...string) (*AuthInterceptor, error) {
+func NewJWTProfileAuthenticator(issuer string, jwtProfileTokenSource JWTProfileTokenSource, scopes ...string) (*AuthInterceptor, error) {
 	ts, err := jwtProfileTokenSource(issuer, scopes)
 	if err != nil {
 		return nil, err
@@ -61,7 +62,7 @@ func NewAuthenticator(issuer string, jwtProfileTokenSource JWTProfileTokenSource
 //
 // Deprecated: use NewAuthenticator(issuer, JWTProfileFromPath(keyPath), scopes...) instead
 func NewAuthInterceptor(issuer, keyPath string, scopes ...string) (*AuthInterceptor, error) {
-	return NewAuthenticator(issuer, JWTProfileFromPath(keyPath), scopes...)
+	return NewJWTProfileAuthenticator(issuer, JWTProfileFromPath(keyPath), scopes...)
 }
 
 func (interceptor *AuthInterceptor) Unary() grpc.UnaryClientInterceptor {
@@ -91,3 +92,15 @@ func (interceptor *AuthInterceptor) setToken(ctx context.Context) (context.Conte
 	}
 	return metadata.AppendToOutgoingContext(ctx, "authorization", token.TokenType+" "+token.AccessToken), nil
 }
+
+type StaticTokenSource string
+
+func (a StaticTokenSource) GetRequestMetadata(ctx context.Context, url ...string) (map[string]string, error) {
+	return map[string]string{
+		"authorization": fmt.Sprintf("Bearer %s", a),
+	}, nil
+}
+
+func (a StaticTokenSource) RequireTransportSecurity() bool {
+	return false
+}
diff --git a/pkg/client/zitadel/client.go b/pkg/client/zitadel/client.go
index 40b3397..6224181 100644
--- a/pkg/client/zitadel/client.go
+++ b/pkg/client/zitadel/client.go
@@ -16,11 +16,13 @@ type Connection struct {
 	issuer                string
 	api                   string
 	jwtProfileTokenSource middleware.JWTProfileTokenSource
+	staticTokenSource     middleware.StaticTokenSource
 	scopes                []string
 	orgID                 string
 	insecure              bool
 	unaryInterceptors     []grpc.UnaryClientInterceptor
 	streamInterceptors    []grpc.StreamClientInterceptor
+	perRPCCredentials     credentials.PerRPCCredentials
 	*grpc.ClientConn
 }
 
@@ -38,17 +40,27 @@ func NewConnection(issuer, api string, scopes []string, options ...Option) (*Con
 		}
 	}
 
-	err := c.setInterceptors(c.issuer, c.orgID, c.scopes, c.jwtProfileTokenSource)
-	if err != nil {
-		return nil, err
-	}
-	dialOptions := []grpc.DialOption{
-		grpc.WithChainUnaryInterceptor(
-			c.unaryInterceptors...,
-		),
-		grpc.WithChainStreamInterceptor(
-			c.streamInterceptors...,
-		),
+	dialOptions := []grpc.DialOption{}
+	if c.staticTokenSource == "" {
+		err := c.setInterceptors(c.issuer, c.orgID, c.scopes, c.jwtProfileTokenSource)
+		if err != nil {
+			return nil, err
+		}
+		dialOptions = append(dialOptions,
+			grpc.WithChainUnaryInterceptor(
+				c.unaryInterceptors...,
+			),
+			grpc.WithChainStreamInterceptor(
+				c.streamInterceptors...,
+			),
+		)
+	} else {
+		c.setCredentials(c.staticTokenSource)
+		dialOptions = append(dialOptions,
+			grpc.WithPerRPCCredentials(
+				c.perRPCCredentials,
+			),
+		)
 	}
 
 	opt, err := transportOption(c.api, c.insecure)
@@ -66,7 +78,7 @@ func NewConnection(issuer, api string, scopes []string, options ...Option) (*Con
 }
 
 func (c *Connection) setInterceptors(issuer, orgID string, scopes []string, jwtProfileTokenSource middleware.JWTProfileTokenSource) error {
-	auth, err := middleware.NewAuthenticator(issuer, jwtProfileTokenSource, scopes...)
+	auth, err := middleware.NewJWTProfileAuthenticator(issuer, jwtProfileTokenSource, scopes...)
 	if err != nil {
 		return err
 	}
@@ -81,6 +93,10 @@ func (c *Connection) setInterceptors(issuer, orgID string, scopes []string, jwtP
 	return nil
 }
 
+func (c *Connection) setCredentials(staticTokenSource middleware.StaticTokenSource) {
+	c.perRPCCredentials = staticTokenSource
+}
+
 func transportOption(api string, insecure bool) (grpc.DialOption, error) {
 	if insecure {
 		return grpc.WithInsecure(), nil
@@ -138,6 +154,15 @@ func WithJWTProfileTokenSource(provider middleware.JWTProfileTokenSource) func(*
 	}
 }
 
+//WithStaticTokenSource sets the provider used for the authentication
+//if not set, the key file will be read from the path set in env var ZITADEL_KEY_PATH
+func WithStaticTokenSource(provider middleware.StaticTokenSource) func(*Connection) error {
+	return func(client *Connection) error {
+		client.staticTokenSource = provider
+		return nil
+	}
+}
+
 //WithOrgID sets the organization context (where the api calls are executed)
 //if not set the resource owner (organisation) of the calling user will be used
 func WithOrgID(orgID string) func(*Connection) error {