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

go.sum

@@ -77,6 +77,7 @@ github.com/gogo/protobuf v0.0.0-20171007142547-342cbe0a0415/go.mod h1:r8qH/GZQm5
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.1 h1:/s5zKNz0uPFCZ5hddgPdo2TK2TVrUNMn0OOX8/aZMTE=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
+github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/groupcache v0.0.0-20160516000752-02826c3e7903/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20190129154638-5b532d6fd5ef/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
@@ -115,6 +116,7 @@ github.com/gophercloud/gophercloud v0.0.0-20190126172459-c818fa66e4c8/go.mod h1:
 github.com/gophercloud/gophercloud v0.0.0-20190307220656-fe1ba5ce12dd/go.mod h1:vxM41WHh5uqHVBMZHzuwNOHh8XEoIEcSTewFxm1c5g8=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
 github.com/gregjones/httpcache v0.0.0-20170728041850-787624de3eb7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
+github.com/grpc-ecosystem/go-grpc-middleware v1.0.0 h1:Iju5GlWwrvL6UBg4zJJt3btmonfrMlCDdsejg4CZE7c=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.8.3/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=

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 {