added server call handler for propagating the authentication to the s…

…pring security context
LogNet · Jul 18, 2021 · b29dc60 · b29dc60
1 parent da062e1
commit b29dc60
Showing 1 changed file with 43 additions and 1 deletion.
diff --git a/...g-boot-starter/src/main/java/org/lognet/springboot/grpc/security/SecurityInterceptor.java b/...g-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()) {