add support for client unary interceptors in pepclient

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/dnstap/golang-dnstap v0.0.0-20170829151710-2cf77a2b5e11
 	github.com/golang/protobuf v1.3.2
 	github.com/google/uuid v1.1.1
+	github.com/grpc-ecosystem/go-grpc-middleware v1.0.0
 	github.com/grpc-ecosystem/grpc-opentracing v0.0.0-20180507213350-8e809c8a8645
 	github.com/infobloxopen/go-trees v0.0.0-20190313150506-2af4e13f9062
 	github.com/miekg/dns v1.1.15

pep/client.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	ot "github.com/opentracing/opentracing-go"
+	"google.golang.org/grpc"
 )
 
 var (
@@ -83,6 +84,14 @@ type Option func(*options)
 
 const virtualServerAddress = "pdp"
 
+// WithClientUnaryInterceptors returns an Options which appends to
+// the client unary interceptors of the underlying connection to pdp
+func WithClientUnaryInterceptors(interceptors ...grpc.UnaryClientInterceptor) Option {
+	return func(o *options) {
+		o.clientUnaryInterceptors = append([]grpc.UnaryClientInterceptor{}, interceptors...)
+	}
+}
+
 // WithRoundRobinBalancer returns an Option which sets round-robin balancer with given set of servers.
 func WithRoundRobinBalancer(addresses ...string) Option {
 	return func(o *options) {
@@ -209,20 +218,21 @@ const (
 )
 
 type options struct {
-	addresses         []string
-	balancer          int
-	tracer            ot.Tracer
-	maxStreams        int
-	ctx               context.Context
-	connTimeout       time.Duration
-	connStateCb       ConnectionStateNotificationCallback
-	autoRequestSize   bool
-	maxRequestSize    uint32
-	noPool            bool
-	cache             bool
-	cacheTTL          time.Duration
-	cacheMaxSize      int
-	onCacheHitHandler OnCacheHitHandler
+	addresses               []string
+	balancer                int
+	tracer                  ot.Tracer
+	maxStreams              int
+	ctx                     context.Context
+	connTimeout             time.Duration
+	connStateCb             ConnectionStateNotificationCallback
+	autoRequestSize         bool
+	maxRequestSize          uint32
+	noPool                  bool
+	cache                   bool
+	cacheTTL                time.Duration
+	cacheMaxSize            int
+	onCacheHitHandler       OnCacheHitHandler
+	clientUnaryInterceptors []grpc.UnaryClientInterceptor
 }
 
 // NewClient creates client instance using given options.

pep/client_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	ot "github.com/opentracing/opentracing-go"
+	"google.golang.org/grpc"
 )
 
 func TestNewClient(t *testing.T) {
@@ -59,6 +60,23 @@ func TestNewClientWithTracer(t *testing.T) {
 	}
 }
 
+var noOpClientInterceptor grpc.UnaryClientInterceptor = func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	return nil
+}
+
+func TestNewClientWithInterceptor(t *testing.T) {
+	ci := noOpClientInterceptor
+	c := NewClient(WithClientUnaryInterceptors(ci))
+	uc, ok := c.(*unaryClient)
+	if !ok {
+		t.Fatalf("Expected *unaryClient from NewClient got %#v", c)
+	}
+
+	if len(uc.opts.clientUnaryInterceptors) == 0 {
+		t.Errorf("Expected noOpClientInterceptor as client option but got %v", uc.opts.clientUnaryInterceptors)
+	}
+}
+
 func TestNewClientWithAutoRequestSize(t *testing.T) {
 	c := NewClient(WithAutoRequestSize(true))
 	uc, ok := c.(*unaryClient)

pep/unary_client.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/allegro/bigcache"
+	"github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	ot "github.com/opentracing/opentracing-go"
 	"google.golang.org/grpc"
@@ -64,20 +65,26 @@ func (c *unaryClient) Connect(addr string) error {
 		}
 	}
 
+	var interceptors []grpc.UnaryClientInterceptor
+
 	if c.opts.tracer != nil {
-		opts = append(opts,
-			grpc.WithUnaryInterceptor(
-				otgrpc.OpenTracingClientInterceptor(
-					c.opts.tracer,
-					otgrpc.IncludingSpans(
-						func(parentSpanCtx ot.SpanContext, method string, req, resp interface{}) bool {
-							return parentSpanCtx != nil
-						},
-					),
+		interceptors = append(interceptors,
+			otgrpc.OpenTracingClientInterceptor(
+				c.opts.tracer,
+				otgrpc.IncludingSpans(
+					func(parentSpanCtx ot.SpanContext, method string, req, resp interface{}) bool {
+						return parentSpanCtx != nil
+					},
 				),
-			),
-		)
+			))
+	}
+
+	if c.opts.clientUnaryInterceptors != nil {
+		interceptors = append(interceptors, c.opts.clientUnaryInterceptors...)
 	}
+	opts = append(opts, grpc.WithUnaryInterceptor(
+		grpc_middleware.ChainUnaryClient(interceptors...),
+	))
 
 	cache, err := newCacheFromOptions(c.opts)
 	if err != nil {