Merge pull request #147 from yandex/arcadia

- TLS support for grpc gun
- Authority pseudo-header for grpc gun

components/grpc/core.go

@@ -2,13 +2,15 @@ package grpc
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"log"
 	"time"
 
 	"github.com/jhump/protoreflect/grpcreflect"
 	"github.com/yandex/pandora/core/warmup"
+	"google.golang.org/grpc/credentials"
 	reflectpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
 
 	"github.com/yandex/pandora/core"
@@ -36,9 +38,15 @@ type Sample struct {
 	ShootTimeSeconds float64
 }
 
+type grpcDialOptions struct {
+	Authority string `config:"authority"`
+}
+
 type GunConfig struct {
-	Target  string        `validate:"required"`
-	Timeout time.Duration `config:"timeout"` // grpc request timeout
+	Target      string          `validate:"required"`
+	Timeout     time.Duration   `config:"timeout"` // grpc request timeout
+	TLS         bool            `config:"tls"`
+	DialOptions grpcDialOptions `config:"dial_options"`
 }
 
 type Gun struct {
@@ -53,11 +61,7 @@ type Gun struct {
 }
 
 func (g *Gun) WarmUp(opts *warmup.Options) (interface{}, error) {
-	conn, err := grpc.Dial(
-		g.conf.Target,
-		grpc.WithInsecure(),
-		grpc.WithTimeout(time.Second),
-		grpc.WithUserAgent("load test, pandora universal grpc shooter"))
+	conn, err := makeGRPCConnect(g.conf.Target, g.conf.TLS, g.conf.DialOptions)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to target: %w", err)
 	}
@@ -101,11 +105,7 @@ func NewGun(conf GunConfig) *Gun {
 }
 
 func (g *Gun) Bind(aggr core.Aggregator, deps core.GunDeps) error {
-	conn, err := grpc.Dial(
-		g.conf.Target,
-		grpc.WithInsecure(),
-		grpc.WithTimeout(time.Second),
-		grpc.WithUserAgent("load test, pandora universal grpc shooter"))
+	conn, err := makeGRPCConnect(g.conf.Target, g.conf.TLS, g.conf.DialOptions)
 	if err != nil {
 		log.Fatalf("FATAL: grpc.Dial failed\n %s\n", err)
 	}
@@ -147,7 +147,6 @@ func (g *Gun) shoot(ammo *Ammo) {
 		log.Fatalf("FATAL: Payload parsing error %s\n", err)
 		return
 	}
-
 	md := method.GetInputType()
 	message := dynamic.NewMessage(md)
 	err = message.UnmarshalJSON(payloadJSON)
@@ -186,6 +185,24 @@ func (g *Gun) shoot(ammo *Ammo) {
 
 }
 
+func makeGRPCConnect(target string, isTLS bool, dialOptions grpcDialOptions) (conn *grpc.ClientConn, err error) {
+	opts := []grpc.DialOption{}
+	if isTLS {
+		opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})))
+	} else {
+		opts = append(opts, grpc.WithInsecure())
+	}
+	opts = append(opts, grpc.WithTimeout(time.Second))
+	opts = append(opts, grpc.WithUserAgent("load test, pandora universal grpc shooter"))
+
+	if dialOptions.Authority != "" {
+		opts = append(opts, grpc.WithAuthority(dialOptions.Authority))
+	}
+
+	conn, err = grpc.Dial(target, opts...)
+	return
+}
+
 func convertGrpcStatus(err error) int {
 	s := status.Convert(err)