From bb5d3b7257725695037f64b1f24e1d15a8ca3a80 Mon Sep 17 00:00:00 2001
From: Niels de Vos <ndevos@redhat.com>
Date: Tue, 21 Dec 2021 18:54:13 +0100
Subject: [PATCH] cleanup: refactor gRPC middleware into
 NewMiddlewareServerOption

Signed-off-by: Niels de Vos <ndevos@redhat.com>
---
 internal/csi-common/server.go |  7 +------
 internal/csi-common/utils.go  | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/internal/csi-common/server.go b/internal/csi-common/server.go
index 2ab4ebc1e..815044b28 100644
--- a/internal/csi-common/server.go
+++ b/internal/csi-common/server.go
@@ -27,7 +27,6 @@ import (
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/csi-addons/spec/lib/go/replication"
-	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/prometheus/client_golang/prometheus"
 	"google.golang.org/grpc"
@@ -104,12 +103,8 @@ func (s *nonBlockingGRPCServer) serve(endpoint, hstOptions string, srv Servers,
 		klog.Fatalf("Failed to listen: %v", err)
 	}
 
-	middleWare := []grpc.UnaryServerInterceptor{contextIDInjector, logGRPC, panicHandler}
-	if metrics {
-		middleWare = append(middleWare, grpc_prometheus.UnaryServerInterceptor)
-	}
 	opts := []grpc.ServerOption{
-		grpc_middleware.WithUnaryServerChain(middleWare...),
+		NewMiddlewareServerOption(metrics),
 	}
 
 	server := grpc.NewServer(opts...)
diff --git a/internal/csi-common/utils.go b/internal/csi-common/utils.go
index 4f0132fde..f8096406d 100644
--- a/internal/csi-common/utils.go
+++ b/internal/csi-common/utils.go
@@ -30,6 +30,8 @@ import (
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	rp "github.com/csi-addons/replication-lib-utils/protosanitizer"
 	"github.com/csi-addons/spec/lib/go/replication"
+	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/kubernetes-csi/csi-lib-utils/protosanitizer"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -105,6 +107,18 @@ func isReplicationRequest(req interface{}) bool {
 	return isReplicationRequest
 }
 
+// NewMiddlewareServerOption creates a new grpc.ServerOption that configures a
+// common format for log messages and other gRPC related handlers.
+func NewMiddlewareServerOption(withMetrics bool) grpc.ServerOption {
+	middleWare := []grpc.UnaryServerInterceptor{contextIDInjector, logGRPC, panicHandler}
+
+	if withMetrics {
+		middleWare = append(middleWare, grpc_prometheus.UnaryServerInterceptor)
+	}
+
+	return grpc_middleware.WithUnaryServerChain(middleWare...)
+}
+
 func getReqID(req interface{}) string {
 	// if req is nil empty string will be returned
 	reqID := ""