From faa31962c2acac95884763dd82acaf036d3402c2 Mon Sep 17 00:00:00 2001
From: zhuojie <zhuojie@checkr.com>
Date: Tue, 6 Oct 2020 13:49:43 -0700
Subject: [PATCH] Add grpc header support

---
 plugins/dsrhub_grpc/dsrhub_grpc.go      | 20 ++++++++++++--------
 plugins/dsrhub_grpc/dsrhub_grpc_test.go |  6 +++---
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/plugins/dsrhub_grpc/dsrhub_grpc.go b/plugins/dsrhub_grpc/dsrhub_grpc.go
index 792da5f..b206a3e 100644
--- a/plugins/dsrhub_grpc/dsrhub_grpc.go
+++ b/plugins/dsrhub_grpc/dsrhub_grpc.go
@@ -5,9 +5,11 @@ import (
 	"fmt"
 	"time"
 
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+
 	"github.com/dsrhub/dsrhub/idl_dsrhub"
 	"github.com/ovh/utask/pkg/plugins/taskplugin"
-	"google.golang.org/grpc"
 )
 
 var (
@@ -15,7 +17,7 @@ var (
 	// nolint
 	Plugin = taskplugin.New(
 		"dsrhub_grpc", "1.0", exec,
-		taskplugin.WithConfig(validConfig, DsrHubGRPCConfig{}),
+		taskplugin.WithConfig(validConfig, DSRHubGRPCConfig{}),
 	)
 )
 
@@ -23,15 +25,16 @@ const (
 	defaultTimeoutSeconds = 10
 )
 
-// DsrHubGRPCConfig is the configuration needed to perform an gRPC client side call
-type DsrHubGRPCConfig struct {
+// DSRHubGRPCConfig is the configuration needed to perform an gRPC client side call
+type DSRHubGRPCConfig struct {
 	URL     string                      `json:"url"`
 	Request idl_dsrhub.CreateDSRRequest `json:"request"`
 	Timeout int                         `json:"timeout,omitempty"` // timeout in seconds
+	Header  map[string]string           `json:"header,omitempty"`
 }
 
 func validConfig(config interface{}) error {
-	cfg := config.(*DsrHubGRPCConfig)
+	cfg := config.(*DSRHubGRPCConfig)
 	if cfg.URL == "" {
 		return fmt.Errorf("invalid dsrhub_grpc config url: empty url")
 	}
@@ -50,8 +53,8 @@ func validConfig(config interface{}) error {
 	return nil
 }
 
-func exec(stepName string, config interface{}, execCtx interface{}) (output interface{}, metadata interface{}, err error) {
-	cfg := config.(*DsrHubGRPCConfig)
+func exec(stepName string, config interface{}, execCtx interface{}) (output interface{}, _ interface{}, err error) {
+	cfg := config.(*DSRHubGRPCConfig)
 
 	// TODO: support secure connection with configuration
 	conn, err := grpc.Dial(cfg.URL, grpc.WithInsecure())
@@ -69,7 +72,8 @@ func exec(stepName string, config interface{}, execCtx interface{}) (output inte
 	)
 	defer cancel()
 
-	res, err := idl_dsrhub.NewDSRHubServiceClient(conn).CreateDSR(ctx, &cfg.Request)
+	header := metadata.New(cfg.Header)
+	res, err := idl_dsrhub.NewDSRHubServiceClient(conn).CreateDSR(ctx, &cfg.Request, grpc.Header(&header))
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/plugins/dsrhub_grpc/dsrhub_grpc_test.go b/plugins/dsrhub_grpc/dsrhub_grpc_test.go
index c6e6092..c26054d 100644
--- a/plugins/dsrhub_grpc/dsrhub_grpc_test.go
+++ b/plugins/dsrhub_grpc/dsrhub_grpc_test.go
@@ -8,12 +8,12 @@ import (
 )
 
 func TestValidConfig(t *testing.T) {
-	var cfg DsrHubGRPCConfig
+	var cfg DSRHubGRPCConfig
 
-	cfg = DsrHubGRPCConfig{}
+	cfg = DSRHubGRPCConfig{}
 	assert.Error(t, validConfig(&cfg))
 
-	cfg = DsrHubGRPCConfig{
+	cfg = DSRHubGRPCConfig{
 		URL: "localhost:50051",
 		Request: idl_dsrhub.CreateDSRRequest{
 			Regulation:         "gdpr",