From cdd6ce9a95a9dff24334c61b9b8207d53d4c76ac Mon Sep 17 00:00:00 2001
From: Dave Syer <dsyer@vmware.com>
Date: Thu, 23 Jan 2025 13:02:37 +0000
Subject: [PATCH] Catch exceptions in interceptors as well

---
 .../sample/GrpcServerIntegrationTests.java    | 41 ++++++++++++++++++-
 .../GrpcExceptionHandlerInterceptor.java      | 22 ++++++----
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/samples/grpc-server/src/test/java/org/springframework/grpc/sample/GrpcServerIntegrationTests.java b/samples/grpc-server/src/test/java/org/springframework/grpc/sample/GrpcServerIntegrationTests.java
index f81c95c..7d6c89d 100644
--- a/samples/grpc-server/src/test/java/org/springframework/grpc/sample/GrpcServerIntegrationTests.java
+++ b/samples/grpc-server/src/test/java/org/springframework/grpc/sample/GrpcServerIntegrationTests.java
@@ -94,7 +94,46 @@ class GrpcServerIntegrationTests {
 	@Nested
 	@SpringBootTest
 	@AutoConfigureInProcessTransport
-	class ServerWithExceptionInInterceptor {
+	class ServerWithExceptionInInterceptorCall {
+
+		@Test
+		void specificErrorResponse(@Autowired GrpcChannelFactory channels) {
+			SimpleGrpc.SimpleBlockingStub client = SimpleGrpc.newBlockingStub(channels.createChannel("0.0.0.0:0"));
+			assertThat(assertThrows(StatusRuntimeException.class,
+					() -> client.sayHello(HelloRequest.newBuilder().setName("foo").build()))
+				.getStatus()
+				.getCode()).isEqualTo(Code.INVALID_ARGUMENT);
+		}
+
+		@TestConfiguration
+		static class TestConfig {
+
+			@Bean
+			@GlobalServerInterceptor
+			public ServerInterceptor exceptionInterceptor() {
+				return new CustomInterceptor();
+			}
+
+			static class CustomInterceptor implements ServerInterceptor {
+
+				@Override
+				public <ReqT, RespT> io.grpc.ServerCall.Listener<ReqT> interceptCall(
+						io.grpc.ServerCall<ReqT, RespT> call, io.grpc.Metadata headers,
+						io.grpc.ServerCallHandler<ReqT, RespT> next) {
+					throw new IllegalArgumentException("test");
+
+				}
+
+			}
+
+		}
+
+	}
+
+	@Nested
+	@SpringBootTest
+	@AutoConfigureInProcessTransport
+	class ServerWithExceptionInInterceptorListener {
 
 		@Test
 		void specificErrorResponse(@Autowired GrpcChannelFactory channels) {
diff --git a/spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandlerInterceptor.java b/spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandlerInterceptor.java
index e73c7d4..0826dc1 100644
--- a/spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandlerInterceptor.java
+++ b/spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandlerInterceptor.java
@@ -62,8 +62,21 @@ public class GrpcExceptionHandlerInterceptor implements ServerInterceptor {
 	@Override
 	public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
 			ServerCallHandler<ReqT, RespT> next) {
-		return new ExceptionHandlerListener<>(next.startCall(call, headers), call,
-				new FallbackHandler(this.exceptionHandler));
+		Listener<ReqT> listener;
+		try {
+			listener = next.startCall(call, headers);
+		}
+		catch (Throwable t) {
+			call.close(this.exceptionHandler.handleException(t), headers(t));
+			listener = new Listener<ReqT>() {
+			};
+		}
+		return new ExceptionHandlerListener<>(listener, call, new FallbackHandler(this.exceptionHandler));
+	}
+
+	private static Metadata headers(Throwable t) {
+		Metadata result = Status.trailersFromThrowable(t);
+		return result != null ? result : new Metadata();
 	}
 
 	static class ExceptionHandlerListener<ReqT, RespT> extends SimpleForwardingServerCallListener<ReqT> {
@@ -109,11 +122,6 @@ public class GrpcExceptionHandlerInterceptor implements ServerInterceptor {
 			}
 		}
 
-		private Metadata headers(Throwable t) {
-			Metadata result = Status.trailersFromThrowable(t);
-			return result != null ? result : new Metadata();
-		}
-
 	}
 
 	static class FallbackHandler implements GrpcExceptionHandler {