Skip to content

Commit 22eb624

Browse files
dsyerdsyer
committed
Enhance context switching
Support for gRPC Context access as well as Kotlin coroutine access to the SecurityContext.
1 parent 4a9ac12 commit 22eb624

File tree

9 files changed

+226
-116
lines changed

9 files changed

+226
-116
lines changed

‎pom.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@
8484
<spring-boot.version>4.0.1</spring-boot.version>
8585
<jackson.version>2.20.1</jackson.version>
8686
<junit.version>6.0.1</junit.version>
87+
<kotlin.coroutines.version>1.10.2</kotlin.coroutines.version>
8788
<assertj.version>3.27.6</assertj.version>
8889
<awaitility.version>4.3.0</awaitility.version>
8990
<mockito.version>5.20.0</mockito.version>

‎spring-grpc-core/pom.xml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,18 @@
7979
<groupId>io.grpc</groupId>
8080
<artifactId>grpc-kotlin-stub</artifactId>
8181
<optional>true</optional>
82+
<exclusions>
83+
<exclusion>
84+
<groupId>javax.annotation</groupId>
85+
<artifactId>javax.annotation-api</artifactId>
86+
</exclusion>
87+
</exclusions>
88+
</dependency>
89+
<dependency>
90+
<groupId>org.jetbrains.kotlinx</groupId>
91+
<artifactId>kotlinx-coroutines-core</artifactId>
92+
<optional>true</optional>
93+
<version>${kotlin.coroutines.version}</version>
8294
</dependency>
8395
<dependency>
8496
<groupId>com.salesforce.servicelibs</groupId>

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/AuthenticationProcessInterceptor.java‎

Lines changed: 5 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@
2828
import org.springframework.security.core.context.SecurityContext;
2929
import org.springframework.security.core.context.SecurityContextHolder;
3030

31-
import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener;
31+
import io.grpc.Context;
32+
import io.grpc.Contexts;
3233
import io.grpc.Metadata;
3334
import io.grpc.ServerCall;
3435
import io.grpc.ServerCall.Listener;
@@ -97,67 +98,13 @@ else if (user == null || !user.isAuthenticated()) {
9798

9899
SecurityContext currentContext = SecurityContextHolder.getContext();
99100
try {
100-
return new SecurityContextClearingListener<>(next.startCall(call, headers), currentContext);
101+
Context context = Context.current().withValue(GrpcSecurity.SECURITY_CONTEXT_KEY, currentContext);
102+
return new SecurityContextHandlerListener<ReqT, RespT>(Contexts.interceptCall(context, call, headers, next),
103+
currentContext);
101104
}
102105
finally {
103106
SecurityContextHolder.clearContext();
104107
}
105108
}
106109

107-
static class SecurityContextClearingListener<ReqT> extends SimpleForwardingServerCallListener<ReqT> {
108-
109-
private final SecurityContext securityContext;
110-
111-
SecurityContextClearingListener(ServerCall.Listener<ReqT> delegate, SecurityContext securityContext) {
112-
super(delegate);
113-
this.securityContext = securityContext;
114-
}
115-
116-
@Override
117-
public void onMessage(ReqT message) {
118-
SecurityContextHolder.setContext(this.securityContext);
119-
try {
120-
super.onMessage(message);
121-
}
122-
finally {
123-
SecurityContextHolder.clearContext();
124-
}
125-
}
126-
127-
@Override
128-
public void onHalfClose() {
129-
SecurityContextHolder.setContext(this.securityContext);
130-
try {
131-
super.onHalfClose();
132-
}
133-
finally {
134-
SecurityContextHolder.clearContext();
135-
}
136-
}
137-
138-
@Override
139-
public void onReady() {
140-
SecurityContextHolder.setContext(this.securityContext);
141-
try {
142-
super.onReady();
143-
}
144-
finally {
145-
SecurityContextHolder.clearContext();
146-
}
147-
}
148-
149-
@Override
150-
public void onCancel() {
151-
super.onCancel();
152-
SecurityContextHolder.clearContext();
153-
}
154-
155-
@Override
156-
public void onComplete() {
157-
super.onComplete();
158-
SecurityContextHolder.clearContext();
159-
}
160-
161-
}
162-
163110
}

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/CoroutineSecurityContextInterceptor.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2025-present the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.grpc.server.security;
18+
19+
import org.springframework.security.core.context.SecurityContext;
20+
import org.springframework.security.core.context.SecurityContextHolder;
21+
22+
import io.grpc.Metadata;
23+
import io.grpc.ServerCall;
24+
import io.grpc.kotlin.CoroutineContextServerInterceptor;
25+
import kotlin.coroutines.CoroutineContext;
26+
27+
/**
28+
* A gRPC server interceptor that integrates Spring Security's {@link SecurityContext}
29+
* with Kotlin coroutines by adding an element to the coroutine context.
30+
*
31+
* @author Dave Syer
32+
*/
33+
public class CoroutineSecurityContextInterceptor extends CoroutineContextServerInterceptor {
34+
35+
@Override
36+
public CoroutineContext coroutineContext(ServerCall<?, ?> call, Metadata metadata) {
37+
return new SecurityContextElement(SecurityContextHolder.getContext());
38+
}
39+
40+
}

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/GrpcSecurity.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,12 @@
3434
import org.springframework.security.config.annotation.SecurityConfigurerAdapter;
3535
import org.springframework.security.config.annotation.authentication.builders.AuthenticationManagerBuilder;
3636
import org.springframework.security.core.Authentication;
37+
import org.springframework.security.core.context.SecurityContext;
3738
import org.springframework.security.core.userdetails.UserDetailsService;
3839
import org.springframework.util.Assert;
3940

4041
import io.grpc.Attributes;
42+
import io.grpc.Context;
4143
import io.grpc.Metadata;
4244
import io.grpc.MethodDescriptor;
4345
import io.micrometer.observation.ObservationRegistry;
@@ -78,6 +80,11 @@ public final class GrpcSecurity
7880
*/
7981
public static final int CONTEXT_FILTER_ORDER = 0;
8082

83+
/**
84+
* Key for the SecurityContext in the gRPC Context.
85+
*/
86+
public static Context.Key<SecurityContext> SECURITY_CONTEXT_KEY = Context.key("spring-security-context");
87+
8188
private @Nullable AuthenticationManager authenticationManager;
8289

8390
private List<GrpcAuthenticationExtractor> authenticationExtractors = new ArrayList<>();

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/SecurityContextElement.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* Copyright 2025-present the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.grpc.server.security;
18+
19+
import java.util.Objects;
20+
21+
import org.springframework.security.core.context.SecurityContext;
22+
import org.springframework.security.core.context.SecurityContextHolder;
23+
24+
import kotlin.coroutines.AbstractCoroutineContextElement;
25+
import kotlin.coroutines.CoroutineContext;
26+
import kotlinx.coroutines.ThreadContextElement;
27+
28+
/**
29+
* Holds a Spring Security {@link SecurityContext} in a Kotlin CoroutineContext.
30+
*
31+
* @author Dave Syer
32+
*/
33+
// To be replaced by official implementation when available from Spring Security
34+
class SecurityContextElement extends AbstractCoroutineContextElement implements ThreadContextElement<SecurityContext> {
35+
36+
/** The context key. */
37+
private static final Key Key = new Key();
38+
39+
private final SecurityContext securityContext;
40+
41+
SecurityContextElement(SecurityContext securityContext) {
42+
super(Key);
43+
this.securityContext = securityContext;
44+
}
45+
46+
@Override
47+
public SecurityContext updateThreadContext(CoroutineContext coroutineContext) {
48+
Objects.requireNonNull(coroutineContext, "coroutineContext must not be null");
49+
SecurityContextHolder.setContext(this.securityContext);
50+
return this.securityContext;
51+
}
52+
53+
@Override
54+
public void restoreThreadContext(CoroutineContext coroutineContext, SecurityContext securityContext) {
55+
Objects.requireNonNull(coroutineContext, "coroutineContext must not be null");
56+
SecurityContextHolder.clearContext();
57+
}
58+
59+
public static final class Key implements CoroutineContext.Key<SecurityContextElement> {
60+
61+
}
62+
63+
}

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/SecurityContextHandlerListener.java‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* Copyright 2024-present the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.grpc.server.security;
18+
19+
import org.springframework.security.core.context.SecurityContext;
20+
import org.springframework.security.core.context.SecurityContextHolder;
21+
22+
import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener;
23+
import io.grpc.ServerCall;
24+
25+
class SecurityContextHandlerListener<ReqT, RespT> extends SimpleForwardingServerCallListener<ReqT> {
26+
27+
private SecurityContext securityContext;
28+
29+
SecurityContextHandlerListener(ServerCall.Listener<ReqT> delegate, SecurityContext securityContext) {
30+
super(delegate);
31+
this.securityContext = securityContext;
32+
}
33+
34+
@Override
35+
public void onMessage(ReqT message) {
36+
SecurityContextHolder.setContext(this.securityContext);
37+
try {
38+
super.onMessage(message);
39+
}
40+
finally {
41+
SecurityContextHolder.clearContext();
42+
}
43+
}
44+
45+
@Override
46+
public void onHalfClose() {
47+
SecurityContextHolder.setContext(this.securityContext);
48+
try {
49+
super.onHalfClose();
50+
}
51+
finally {
52+
SecurityContextHolder.clearContext();
53+
}
54+
}
55+
56+
@Override
57+
public void onReady() {
58+
SecurityContextHolder.setContext(this.securityContext);
59+
try {
60+
super.onReady();
61+
}
62+
finally {
63+
SecurityContextHolder.clearContext();
64+
}
65+
}
66+
67+
@Override
68+
public void onCancel() {
69+
super.onCancel();
70+
SecurityContextHolder.clearContext();
71+
}
72+
73+
@Override
74+
public void onComplete() {
75+
super.onComplete();
76+
SecurityContextHolder.clearContext();
77+
}
78+
79+
}

‎spring-grpc-core/src/main/java/org/springframework/grpc/server/security/SecurityContextServerInterceptor.java‎

Lines changed: 5 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@
2020
import org.springframework.security.core.context.SecurityContext;
2121
import org.springframework.security.core.context.SecurityContextHolder;
2222

23-
import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener;
23+
import io.grpc.Context;
24+
import io.grpc.Contexts;
2425
import io.grpc.Metadata;
2526
import io.grpc.ServerCall;
2627
import 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

Comments
 (0)