Move metadata validator to logrus interceptor

dynastymasra · Feb 14, 2020 · 044eced · 044eced
1 parent f4c7ee4
commit 044eced
Showing 1 changed file with 16 additions and 34 deletions.
diff --git a/grpc/interceptor.go b/grpc/interceptor.go
@@ -37,48 +37,23 @@ func AuthInterceptor(expectedToken string) grpc_auth.AuthFunc {
 	}
 }
 
-//ValidateMetadataKeyUnaryInterceptor Function used to check metadata key included in request
-func ValidateMetadataKeyUnaryInterceptor(keys ...string) grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-		for _, key := range keys {
-			val := metautils.ExtractIncoming(ctx).Get(key)
-			if val == "" {
-				return nil, status.Errorf(codes.FailedPrecondition, "metadata key %s not found", key)
-			}
-		}
-
-		resp, err := handler(ctx, req)
-
-		return resp, err
-	}
-}
-
-//ValidateMetadataKeyStreamInterceptor Function used to check metadata key included in request
-func ValidateMetadataKeyStreamInterceptor(keys ...string) grpc.StreamServerInterceptor {
-	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		for _, key := range keys {
-			val := metautils.ExtractIncoming(stream.Context()).Get(key)
-			if val == "" {
-				return status.Errorf(codes.FailedPrecondition, "metadata key not found")
-			}
-		}
-
-		err := handler(srv, stream)
-
-		return err
-	}
-}
-
 /**
 //LogrusUnaryInterceptor gRPC interceptor to log unary request duration status
 this function will set field request_id from metadata
 */
-func LogrusUnaryInterceptor(logger *logrus.Entry) grpc.UnaryServerInterceptor {
+func LogrusUnaryInterceptor(logger *logrus.Entry, keys ...string) grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 		startTime := time.Now()
 
 		resp, err := handler(ctx, req)
 
+		for _, key := range keys {
+			val := metautils.ExtractIncoming(ctx).Get(key)
+			if val == "" {
+				return nil, status.Errorf(codes.FailedPrecondition, "metadata key %s not found", key)
+			}
+		}
+
 		log := logger.WithFields(logrus.Fields{
 			"full_method":  info.FullMethod,
 			"request_time": startTime.Format(time.RFC3339),
@@ -111,12 +86,19 @@ func LogrusUnaryInterceptor(logger *logrus.Entry) grpc.UnaryServerInterceptor {
 LogrusStreamInterceptor gRPC interceptor to log stream request duration status
 this function will set field request_id from metadata
 */
-func LogrusStreamInterceptor(logger *logrus.Entry) grpc.StreamServerInterceptor {
+func LogrusStreamInterceptor(logger *logrus.Entry, keys ...string) grpc.StreamServerInterceptor {
 	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		startTime := time.Now()
 
 		err := handler(srv, stream)
 
+		for _, key := range keys {
+			val := metautils.ExtractIncoming(stream.Context()).Get(key)
+			if val == "" {
+				return status.Errorf(codes.FailedPrecondition, "metadata key not found")
+			}
+		}
+
 		log := logger.WithFields(logrus.Fields{
 			"full_method":  info.FullMethod,
 			"request_time": startTime.Format(time.RFC3339),