From b29dc6067236c2be72ab080bf6d960c3f5800ec9 Mon Sep 17 00:00:00 2001
From: Mark Banierink <mark.banierink@nedap.com>
Date: Sat, 17 Jul 2021 01:05:57 +0200
Subject: [PATCH] added server call handler for propagating the authentication
 to the spring security context

---
 .../grpc/security/SecurityInterceptor.java    | 44 ++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/security/SecurityInterceptor.java b/grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/security/SecurityInterceptor.java
index 41af4646..eca6246f 100644
--- a/grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/security/SecurityInterceptor.java
+++ b/grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/security/SecurityInterceptor.java
@@ -95,7 +95,7 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
             Context ctx = Context.current()
                     .withValue(GrpcSecurity.AUTHENTICATION_CONTEXT_KEY, SecurityContextHolder.getContext().getAuthentication());
 
-            return Contexts.interceptCall(ctx, call, headers, next);
+            return Contexts.interceptCall(ctx, call, headers, authenticationPropagatingHandler(next));
         } catch (AccessDeniedException e) {
             return fail(next, call, headers, Status.PERMISSION_DENIED, e);
         } catch (Exception e) {
@@ -107,6 +107,48 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
 
     }
 
+    private <ReqT, RespT> ServerCallHandler<ReqT, RespT> authenticationPropagatingHandler(ServerCallHandler<ReqT, RespT> next) {
+
+        return (call, headers) -> new ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT>(next.startCall(call, headers)) {
+
+            @Override
+            public void onMessage(ReqT message) {
+                propagateAuthentication(() -> super.onMessage(message));
+            }
+
+            @Override
+            public void onHalfClose() {
+                propagateAuthentication(super::onHalfClose);
+            }
+
+            @Override
+            public void onCancel() {
+                propagateAuthentication(super::onCancel);
+            }
+
+            @Override
+            public void onComplete() {
+                propagateAuthentication(super::onComplete);
+            }
+
+            @Override
+            public void onReady() {
+                propagateAuthentication(super::onReady);
+            }
+
+            private void propagateAuthentication(Runnable runnable) {
+                try {
+                    SecurityContextHolder.getContext().setAuthentication(GrpcSecurity.AUTHENTICATION_CONTEXT_KEY.get());
+                    runnable.run();
+                } finally {
+                    SecurityContextHolder.clearContext();
+                }
+            }
+
+        };
+
+    }
+
     private <RespT, ReqT> ServerCall.Listener<ReqT> fail(ServerCallHandler<ReqT, RespT> next, ServerCall<ReqT, RespT> call, Metadata headers,final Status status, Exception exception) {
 
         if (authCfg.isFailFast()) {