pass headers to reflection client

ktr0731 · Jan 28, 2021 · ab5c40e · ab5c40e
1 parent 62389c7
commit ab5c40e
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 13 deletions.
diff --git a/grpc/grpc.go b/grpc/grpc.go
@@ -114,7 +114,7 @@ type client struct {
 // The set of cert and certKey enables mutual authentication if useTLS is enabled.
 // If one of it is not found, NewClient returns ErrMutualAuthParamsAreNotEnough.
 // If useTLS is false, cacert, cert and certKey are ignored.
-func NewClient(addr, serverName string, useReflection, useTLS bool, cacert, cert, certKey string) (Client, error) {
+func NewClient(addr, serverName string, useReflection, useTLS bool, cacert, cert, certKey string, headers map[string][]string) (Client, error) {
 	var opts []grpc.DialOption
 	if !useTLS {
 		opts = append(opts, grpc.WithInsecure())
@@ -163,7 +163,7 @@ func NewClient(addr, serverName string, useReflection, useTLS bool, cacert, cert
 	}
 
 	if useReflection {
-		client.Client = grpcreflection.NewClient(conn)
+		client.Client = grpcreflection.NewClient(conn, headers)
 	}
 
 	return client, nil

diff --git a/grpc/grpc_test.go b/grpc/grpc_test.go
@@ -46,7 +46,7 @@ func TestNewClient(t *testing.T) {
 	for name, c := range cases {
 		c := c
 		t.Run(name, func(t *testing.T) {
-			_, err := NewClient(c.addr, "", c.useReflection, c.useTLS, c.cacert, c.cert, c.certKey)
+			_, err := NewClient(c.addr, "", c.useReflection, c.useTLS, c.cacert, c.cert, c.certKey, nil)
 			if c.err != nil {
 				if err == nil {
 					t.Fatalf("NewClient must return an error, but got nil")

diff --git a/grpc/grpcreflection/reflection.go b/grpc/grpcreflection/reflection.go
@@ -12,6 +12,7 @@ import (
 	"github.com/ktr0731/grpc-web-go-client/grpcweb/grpcweb_reflection_v1alpha"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
 	"google.golang.org/grpc/status"
 )
@@ -35,17 +36,21 @@ type client struct {
 	client *gr.Client
 }
 
+func getCtx(headers map[string][]string) context.Context {
+	return metadata.NewOutgoingContext(context.Background(), metadata.MD(headers))
+}
+
 // NewClient returns an instance of gRPC reflection client for gRPC protocol.
-func NewClient(conn grpc.ClientConnInterface) Client {
+func NewClient(conn grpc.ClientConnInterface, headers map[string][]string) Client {
 	return &client{
-		client: gr.NewClient(context.Background(), grpc_reflection_v1alpha.NewServerReflectionClient(conn)),
+		client: gr.NewClient(getCtx(headers), grpc_reflection_v1alpha.NewServerReflectionClient(conn)),
 	}
 }
 
 // NewWebClient returns an instance of gRPC reflection client for gRPC-Web protocol.
-func NewWebClient(conn *grpcweb.ClientConn) Client {
+func NewWebClient(conn *grpcweb.ClientConn, headers map[string][]string) Client {
 	return &client{
-		client: gr.NewClient(context.Background(), grpcweb_reflection_v1alpha.NewServerReflectionClient(conn)),
+		client: gr.NewClient(getCtx(headers), grpcweb_reflection_v1alpha.NewServerReflectionClient(conn)),
 	}
 }
 

diff --git a/grpc/web.go b/grpc/web.go
@@ -18,7 +18,7 @@ type webClient struct {
 	grpcreflection.Client
 }
 
-func NewWebClient(addr string, useReflection, useTLS bool, cacert, cert, certKey string) Client {
+func NewWebClient(addr string, useReflection, useTLS bool, cacert, cert, certKey string, headers Headers) Client {
 	conn, err := grpcweb.DialContext(addr)
 	if err != nil {
 		panic(err)
@@ -29,7 +29,7 @@ func NewWebClient(addr string, useReflection, useTLS bool, cacert, cert, certKey
 	}
 
 	if useReflection {
-		client.Client = grpcreflection.NewWebClient(conn)
+		client.Client = grpcreflection.NewWebClient(conn, headers)
 	}
 
 	return client

diff --git a/grpc/web_test.go b/grpc/web_test.go
@@ -8,7 +8,7 @@ import (
 )
 
 func TestWebClient(t *testing.T) {
-	client := grpc.NewWebClient("", false, false, "", "", "")
+	client := grpc.NewWebClient("", false, false, "", "", "", nil)
 	t.Run("Invoke returns an error if FQRN is invalid", func(t *testing.T) {
 		_, _, err := client.Invoke(context.Background(), "invalid-fqrn", nil, nil)
 		if err == nil {

diff --git a/mode/common.go b/mode/common.go
@@ -31,7 +31,7 @@ func newGRPCClient(cfg *config.Config) (grpc.Client, error) {
 	addr := fmt.Sprintf("%s:%s", cfg.Server.Host, cfg.Server.Port)
 	if cfg.Request.Web {
 		//TODO: remove second arg
-		return grpc.NewWebClient(addr, cfg.Server.Reflection, false, "", "", ""), nil
+		return grpc.NewWebClient(addr, cfg.Server.Reflection, false, "", "", "", grpc.Headers(cfg.Request.Header)), nil
 	}
 	client, err := grpc.NewClient(
 		addr,
@@ -40,7 +40,8 @@ func newGRPCClient(cfg *config.Config) (grpc.Client, error) {
 		cfg.Server.TLS,
 		cfg.Request.CACertFile,
 		cfg.Request.CertFile,
-		cfg.Request.CertKeyFile)
+		cfg.Request.CertKeyFile,
+		cfg.Request.Header)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to instantiate a gRPC client")
 	}

diff --git a/usecase/header_test.go b/usecase/header_test.go
@@ -9,7 +9,7 @@ import (
 
 func TestHeader(t *testing.T) {
 	defer Clear()
-	client, err := grpc.NewClient("", "", false, false, "", "", "")
+	client, err := grpc.NewClient("", "", false, false, "", "", "", nil)
 	if err != nil {
 		t.Fatalf("grpc.NewClient must not return an error, but got '%s'", err)
 	}