From 83a3b3b382d465575c93a51326fa5d80d6ae574c Mon Sep 17 00:00:00 2001
From: Jan Tattermusch <jtattermusch@google.com>
Date: Thu, 31 Jan 2019 11:35:59 +0100
Subject: [PATCH] fix ServerServiceDefinition interception

---
 .../ServerServiceDefinitionExtensions.cs      | 52 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/src/csharp/Grpc.Core/Interceptors/ServerServiceDefinitionExtensions.cs b/src/csharp/Grpc.Core/Interceptors/ServerServiceDefinitionExtensions.cs
index 36eb08cf528..321cf080faa 100644
--- a/src/csharp/Grpc.Core/Interceptors/ServerServiceDefinitionExtensions.cs
+++ b/src/csharp/Grpc.Core/Interceptors/ServerServiceDefinitionExtensions.cs
@@ -44,7 +44,10 @@ namespace Grpc.Core.Interceptors
         {
             GrpcPreconditions.CheckNotNull(serverServiceDefinition, nameof(serverServiceDefinition));
             GrpcPreconditions.CheckNotNull(interceptor, nameof(interceptor));
-            return new ServerServiceDefinition(serverServiceDefinition.GetCallHandlers().ToDictionary(x => x.Key, x => x.Value.Intercept(interceptor)));
+
+            var binder = new InterceptingServiceBinder(interceptor);
+            serverServiceDefinition.BindService(binder);
+            return binder.GetInterceptedServerServiceDefinition();
         }
 
         /// <summary>
@@ -75,5 +78,52 @@ namespace Grpc.Core.Interceptors
 
             return serverServiceDefinition;
         }
+
+        /// <summary>
+        /// Helper for creating <c>ServerServiceDefinition</c> with intercepted handlers.
+        /// </summary>
+        private class InterceptingServiceBinder : ServiceBinderBase
+        {
+            readonly ServerServiceDefinition.Builder builder = ServerServiceDefinition.CreateBuilder();
+            readonly Interceptor interceptor;
+
+            public InterceptingServiceBinder(Interceptor interceptor)
+            {
+                this.interceptor = GrpcPreconditions.CheckNotNull(interceptor, nameof(interceptor));
+            }
+
+            internal ServerServiceDefinition GetInterceptedServerServiceDefinition()
+            {
+                return builder.Build();
+            }
+
+            public override void AddMethod<TRequest, TResponse>(
+                Method<TRequest, TResponse> method,
+                UnaryServerMethod<TRequest, TResponse> handler)
+            {
+                builder.AddMethod(method, (request, context) => interceptor.UnaryServerHandler(request, context, handler));
+            }
+
+            public override void AddMethod<TRequest, TResponse>(
+                Method<TRequest, TResponse> method,
+                ClientStreamingServerMethod<TRequest, TResponse> handler)
+            {
+                builder.AddMethod(method, (requestStream, context) => interceptor.ClientStreamingServerHandler(requestStream, context, handler));
+            }
+
+            public override void AddMethod<TRequest, TResponse>(
+                Method<TRequest, TResponse> method,
+                ServerStreamingServerMethod<TRequest, TResponse> handler)
+            {
+                builder.AddMethod(method, (request, responseStream, context) => interceptor.ServerStreamingServerHandler(request, responseStream, context, handler));
+            }
+
+            public override void AddMethod<TRequest, TResponse>(
+                Method<TRequest, TResponse> method,
+                DuplexStreamingServerMethod<TRequest, TResponse> handler)
+            {
+                builder.AddMethod(method, (requestStream, responseStream, context) => interceptor.DuplexStreamingServerHandler(requestStream, responseStream, context, handler));
+            }
+        }
     }
 }