NETOBSERV-333: allow overriding default GRPC options (#31)

Author: Mario Macias

pkg/grpc/grpc_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/netobserv/netobserv-ebpf-agent/pkg/pbflow"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -84,6 +85,40 @@ func TestGRPCCommunication(t *testing.T) {
 	}
 }
 
+func TestConstructorOptions(t *testing.T) {
+	port, err := test.FreeTCPPort()
+	require.NoError(t, err)
+	intercepted := make(chan struct{})
+	// Override the default GRPC collector to verify that StartCollector is applying the
+	// passed options
+	_, err = StartCollector(port, make(chan *pbflow.Records),
+		WithGRPCServerOptions(grpc.UnaryInterceptor(func(
+			ctx context.Context,
+			req interface{},
+			info *grpc.UnaryServerInfo,
+			handler grpc.UnaryHandler,
+		) (resp interface{}, err error) {
+			close(intercepted)
+			return handler(ctx, req)
+		})))
+	require.NoError(t, err)
+	cc, err := ConnectClient(fmt.Sprintf("127.0.0.1:%d", port))
+	require.NoError(t, err)
+	client := cc.Client()
+
+	go func() {
+		_, err = client.Send(context.Background(),
+			&pbflow.Records{Entries: []*pbflow.Record{{EthProtocol: 123, Bytes: 456}}})
+		require.NoError(t, err)
+	}()
+
+	select {
+	case <-intercepted:
+	case <-time.After(timeout):
+		require.Fail(t, "timeout waiting for unary interceptor to work")
+	}
+}
+
 func BenchmarkGRPCCommunication(b *testing.B) {
 	port, err := test.FreeTCPPort()
 	require.NoError(b, err)

pkg/grpc/server.go

@@ -16,15 +16,35 @@ type CollectorServer struct {
 	grpcServer *grpc.Server
 }
 
+type collectorOptions struct {
+	grpcServerOptions []grpc.ServerOption
+}
+
+// CollectorOption allows overriding the default configuration of the CollectorServer instance.
+// Use them in the StartCollector function.
+type CollectorOption func(options *collectorOptions)
+
+func WithGRPCServerOptions(options ...grpc.ServerOption) CollectorOption {
+	return func(copt *collectorOptions) {
+		copt.grpcServerOptions = options
+	}
+}
+
 // StartCollector listens in background for gRPC+Protobuf flows in the given port, and forwards each
 // set of *pbflow.Records by the provided channel.
-func StartCollector(port int, recordForwarder chan<- *pbflow.Records) (*CollectorServer, error) {
+func StartCollector(
+	port int, recordForwarder chan<- *pbflow.Records, options ...CollectorOption,
+) (*CollectorServer, error) {
+	copts := collectorOptions{}
+	for _, opt := range options {
+		opt(&copts)
+	}
+
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {
 		return nil, err
 	}
-	// TODO: set server option arguments
-	grpcServer := grpc.NewServer()
+	grpcServer := grpc.NewServer(copts.grpcServerOptions...)
 	pbflow.RegisterCollectorServer(grpcServer, &collectorAPI{
 		recordForwarder: recordForwarder,
 	})