2020import org .springframework .security .core .context .SecurityContext ;
2121import org .springframework .security .core .context .SecurityContextHolder ;
2222
23- import io .grpc .ForwardingServerCallListener .SimpleForwardingServerCallListener ;
23+ import io .grpc .Context ;
24+ import io .grpc .Contexts ;
2425import io .grpc .Metadata ;
2526import io .grpc .ServerCall ;
2627import io .grpc .ServerCall .Listener ;
@@ -38,63 +39,9 @@ public int getOrder() {
3839 public <ReqT , RespT > Listener <ReqT > interceptCall (ServerCall <ReqT , RespT > call , Metadata headers ,
3940 ServerCallHandler <ReqT , RespT > next ) {
4041 SecurityContext securityContext = SecurityContextHolder .getContext ();
41- return new SecurityContextHandlerListener <ReqT , RespT >(next .startCall (call , headers ), securityContext );
42- }
43-
44- static class SecurityContextHandlerListener <ReqT , RespT > extends SimpleForwardingServerCallListener <ReqT > {
45-
46- private SecurityContext securityContext ;
47-
48- SecurityContextHandlerListener (ServerCall .Listener <ReqT > delegate , SecurityContext securityContext ) {
49- super (delegate );
50- this .securityContext = securityContext ;
51- }
52-
53- @ Override
54- public void onMessage (ReqT message ) {
55- SecurityContextHolder .setContext (this .securityContext );
56- try {
57- super .onMessage (message );
58- }
59- finally {
60- SecurityContextHolder .clearContext ();
61- }
62- }
63-
64- @ Override
65- public void onHalfClose () {
66- SecurityContextHolder .setContext (this .securityContext );
67- try {
68- super .onHalfClose ();
69- }
70- finally {
71- SecurityContextHolder .clearContext ();
72- }
73- }
74-
75- @ Override
76- public void onReady () {
77- SecurityContextHolder .setContext (this .securityContext );
78- try {
79- super .onReady ();
80- }
81- finally {
82- SecurityContextHolder .clearContext ();
83- }
84- }
85-
86- @ Override
87- public void onCancel () {
88- super .onCancel ();
89- SecurityContextHolder .clearContext ();
90- }
91-
92- @ Override
93- public void onComplete () {
94- super .onComplete ();
95- SecurityContextHolder .clearContext ();
96- }
97-
42+ Context context = Context .current ().withValue (GrpcSecurity .SECURITY_CONTEXT_KEY , securityContext );
43+ return new SecurityContextHandlerListener <ReqT , RespT >(Contexts .interceptCall (context , call , headers , next ),
44+ securityContext );
9845 }
9946
10047}
0 commit comments