Skip to content

Commit 7237313

Browse files
mikereichemikereiche
committed
Propagate SecurityContext into @transactional methods. (#1979)
Closes #1944.
1 parent 72fdb1c commit 7237313

File tree

6 files changed

+147
-7
lines changed

6 files changed

+147
-7
lines changed

‎src/main/java/org/springframework/data/couchbase/transaction/CouchbaseCallbackTransactionManager.java

+38-7
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import reactor.core.publisher.Flux;
1919
import reactor.core.publisher.Mono;
2020

21+
import java.lang.reflect.InvocationTargetException;
2122
import java.time.Duration;
2223
import java.util.ArrayList;
2324
import java.util.List;
@@ -89,6 +90,7 @@ public <T> T execute(TransactionDefinition definition, TransactionCallback<T> ca
8990
@Stability.Internal
9091
<T> Flux<T> executeReactive(TransactionDefinition definition,
9192
org.springframework.transaction.reactive.TransactionCallback<T> callback) {
93+
final CouchbaseResourceHolder couchbaseResourceHolder = new CouchbaseResourceHolder(null, getSecurityContext()); // caller's resources
9294
return TransactionalSupport.checkForTransactionInThreadLocalStorage().flatMapMany(isInTransaction -> {
9395
boolean isInExistingTransaction = isInTransaction.isPresent();
9496
boolean createNewTransaction = handlePropagation(definition, isInExistingTransaction);
@@ -100,17 +102,20 @@ <T> Flux<T> executeReactive(TransactionDefinition definition,
100102
} else {
101103
return Mono.error(new UnsupportedOperationException("Unsupported operation"));
102104
}
103-
});
105+
}).contextWrite( // set CouchbaseResourceHolder containing caller's SecurityContext
106+
ctx -> ctx.put(CouchbaseResourceHolder.class, couchbaseResourceHolder));
104107
}
105108

106109
private <T> T executeNewTransaction(TransactionCallback<T> callback) {
107110
final AtomicReference<T> execResult = new AtomicReference<>();
111+
final CouchbaseResourceHolder couchbaseResourceHolder = new CouchbaseResourceHolder(null, getSecurityContext());
108112

109113
// Each of these transactions will block one thread on the underlying SDK's transactions scheduler. This
110114
// scheduler is effectively unlimited, but this can still potentially lead to high thread usage by the application.
111115
// If this is an issue then users need to instead use the standard Couchbase reactive transactions SDK.
112116
try {
113117
TransactionResult ignored = couchbaseClientFactory.getCluster().transactions().run(ctx -> {
118+
setSecurityContext(couchbaseResourceHolder.getSecurityContext()); // set the security context for the transaction
114119
CouchbaseTransactionStatus status = new CouchbaseTransactionStatus(ctx, true, false, false, true, null);
115120

116121
T res = callback.doInTransaction(status);
@@ -173,12 +178,16 @@ public boolean isCompleted() {
173178
}
174179
};
175180

176-
return Flux.from(callback.doInTransaction(status)).doOnNext(v -> out.add(v)).then(Mono.defer(() -> {
177-
if (status.isRollbackOnly()) {
178-
return Mono.error(new TransactionRollbackRequestedException("TransactionStatus.isRollbackOnly() is set"));
179-
}
180-
return Mono.empty();
181-
}));
181+
// Get caller's resources, set SecurityContext for the transaction
182+
return CouchbaseResourceOwner.get().map(cbrh -> setSecurityContext(cbrh.get().getSecurityContext()))
183+
.flatMap(ignore -> Flux.from(callback.doInTransaction(status)).doOnNext(v -> out.add(v))
184+
.then(Mono.defer(() -> {
185+
if (status.isRollbackOnly()) {
186+
return Mono.error(new TransactionRollbackRequestedException(
187+
"TransactionStatus.isRollbackOnly() is set"));
188+
}
189+
return Mono.empty();
190+
})));
182191
});
183192

184193
}, this.options).thenMany(Flux.defer(() -> Flux.fromIterable(out))).onErrorMap(ex -> {
@@ -288,4 +297,26 @@ public void rollback(TransactionStatus ignored) throws TransactionException {
288297
throw new UnsupportedOperationException(
289298
"Direct programmatic use of the Couchbase PlatformTransactionManager is not supported");
290299
}
300+
301+
static private Object getSecurityContext() {
302+
try {
303+
Class<?> securityContextHolderClass = Class
304+
.forName("org.springframework.security.core.context.SecurityContextHolder");
305+
return securityContextHolderClass.getMethod("getContext").invoke(null);
306+
} catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException
307+
| InvocationTargetException cnfe) {}
308+
return null;
309+
}
310+
311+
static private <S> S setSecurityContext(S sc) {
312+
try {
313+
Class<?> securityContextHolder = Class.forName("org.springframework.security.core.context.SecurityContext");
314+
Class<?> securityContextHolderClass = Class
315+
.forName("org.springframework.security.core.context.SecurityContextHolder");
316+
securityContextHolderClass.getMethod("setContext", new Class[] { securityContextHolder }).invoke(null, sc);
317+
} catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException
318+
| InvocationTargetException cnfe) {}
319+
return sc;
320+
}
321+
291322
}

‎src/main/java/org/springframework/data/couchbase/transaction/CouchbaseResourceHolder.java

+20
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
public class CouchbaseResourceHolder extends ResourceHolderSupport {
3535

3636
private @Nullable CoreTransactionAttemptContext core; // which holds the atr
37+
private @Nullable Object securityContext; // SecurityContext. We don't have the class.
38+
3739
Map<Integer, Object> getResultMap = new HashMap<>();
3840

3941
/**
@@ -42,7 +44,17 @@ public class CouchbaseResourceHolder extends ResourceHolderSupport {
4244
* @param core the associated {@link CoreTransactionAttemptContext}. Can be {@literal null}.
4345
*/
4446
public CouchbaseResourceHolder(@Nullable CoreTransactionAttemptContext core) {
47+
this(core, null);
48+
}
49+
50+
/**
51+
* Create a new {@link CouchbaseResourceHolder} for a given {@link CoreTransactionAttemptContext session}.
52+
*
53+
* @param core the associated {@link CoreTransactionAttemptContext}. Can be {@literal null}.
54+
*/
55+
public CouchbaseResourceHolder(@Nullable CoreTransactionAttemptContext core, @Nullable Object securityContext) {
4556
this.core = core;
57+
this.securityContext = securityContext;
4658
}
4759

4860
/**
@@ -53,6 +65,14 @@ public CoreTransactionAttemptContext getCore() {
5365
return core;
5466
}
5567

68+
/**
69+
* @return the associated {@link CoreTransactionAttemptContext}. Can be {@literal null}.
70+
*/
71+
@Nullable
72+
public Object getSecurityContext() {
73+
return securityContext;
74+
}
75+
5676
public Object transactionResultHolder(Object holder, Object o) {
5777
getResultMap.put(System.identityHashCode(o), holder);
5878
return holder;

‎src/main/java/org/springframework/data/couchbase/transaction/CouchbaseResourceOwner.java

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package org.springframework.data.couchbase.transaction;
2+
3+
import reactor.core.publisher.Mono;
4+
5+
import java.util.Optional;
6+
7+
import com.couchbase.client.core.annotation.Stability.Internal;
8+
9+
@Internal
10+
public class CouchbaseResourceOwner {
11+
private static final ThreadLocal<CouchbaseResourceHolder> marker = new ThreadLocal();
12+
13+
public CouchbaseResourceOwner() {}
14+
15+
public static void set(CouchbaseResourceHolder toInject) {
16+
if (marker.get() != null) {
17+
throw new IllegalStateException(
18+
"Trying to set resource holder when already inside a transaction - likely an internal bug, please report it");
19+
} else {
20+
marker.set(toInject);
21+
}
22+
}
23+
24+
public static void clear() {
25+
marker.remove();
26+
}
27+
28+
public static Mono<Optional<CouchbaseResourceHolder>> get() {
29+
return Mono.deferContextual((ctx) -> {
30+
CouchbaseResourceHolder fromThreadLocal = marker.get();
31+
CouchbaseResourceHolder fromReactive = ctx.hasKey(CouchbaseResourceHolder.class)
32+
? ctx.get(CouchbaseResourceHolder.class)
33+
: null;
34+
if (fromThreadLocal != null) {
35+
return Mono.just(Optional.of(fromThreadLocal));
36+
} else {
37+
return fromReactive != null ? Mono.just(Optional.of(fromReactive)) : Mono.just(Optional.empty());
38+
}
39+
});
40+
}
41+
}

‎src/test/java/org/springframework/data/couchbase/transactions/CouchbasePersonTransactionReactiveIntegrationTests.java

+13
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.couchbase.client.java.query.QueryScanConsistency.REQUEST_PLUS;
2020

21+
import org.springframework.data.couchbase.util.Util;
2122
import reactor.core.publisher.Mono;
2223
import reactor.test.StepVerifier;
2324

@@ -119,6 +120,7 @@ public void shouldRollbackAfterExceptionOfTxAnnotatedMethod() {
119120
@Test
120121
public void commitShouldPersistTxEntries() {
121122

123+
System.err.println("parent SecurityContext: " + System.identityHashCode(Util.getSecurityContext()));
122124
personService.savePerson(WalterWhite) //
123125
.as(StepVerifier::create) //
124126
.expectNextCount(1) //
@@ -130,6 +132,17 @@ public void commitShouldPersistTxEntries() {
130132
.verifyComplete();
131133
}
132134

135+
@Test
136+
public void commitShouldPersistTxEntriesBlocking() {
137+
System.err.println("parent SecurityContext: " + System.identityHashCode(Util.getSecurityContext()));
138+
Person p = personService.savePersonBlocking(WalterWhite);
139+
140+
operations.findByQuery(Person.class).withConsistency(REQUEST_PLUS).count() //
141+
.as(StepVerifier::create) //
142+
.expectNext(1L) //
143+
.verifyComplete();
144+
}
145+
133146
@Test
134147
public void commitShouldPersistTxEntriesOfTxAnnotatedMethod() {
135148

‎src/test/java/org/springframework/data/couchbase/transactions/PersonServiceReactive.java

+12
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121

2222
import org.springframework.data.couchbase.core.TransactionalSupport;
2323
import org.springframework.data.couchbase.domain.PersonWithoutVersion;
24+
import org.springframework.data.couchbase.transaction.CouchbaseResourceHolder;
25+
import org.springframework.data.couchbase.util.Util;
2426
import org.springframework.stereotype.Service;
2527
import reactor.core.publisher.Flux;
2628
import reactor.core.publisher.Mono;
@@ -31,6 +33,9 @@
3133
import org.springframework.transaction.annotation.Transactional;
3234
import org.springframework.transaction.reactive.TransactionalOperator;
3335

36+
import java.lang.reflect.InvocationTargetException;
37+
import java.util.Optional;
38+
3439
/**
3540
* reactive PersonService for tests
3641
*
@@ -57,8 +62,15 @@ public Mono<Person> savePersonErrors(Person person) {
5762
.<Person> flatMap(it -> Mono.error(new SimulateFailureException()));
5863
}
5964

65+
@Transactional
66+
public Person savePersonBlocking(Person person) {
67+
System.err.println("savePerson: "+Thread.currentThread().getName() +" "+ System.identityHashCode(Util.getSecurityContext()));
68+
return personOperations.insertById(Person.class).one(person);
69+
}
70+
6071
@Transactional
6172
public Mono<Person> savePerson(Person person) {
73+
System.err.println("savePerson: "+Thread.currentThread().getName() +" "+ System.identityHashCode(Util.getSecurityContext()));
6274
return TransactionalSupport.checkForTransactionInThreadLocalStorage().map(stat -> {
6375
assertTrue(stat.isPresent(), "Not in transaction");
6476
System.err.println("In a transaction!!");

‎src/test/java/org/springframework/data/couchbase/util/Util.java

+23
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import static org.awaitility.Awaitility.with;
2121

2222
import java.io.InputStream;
23+
import java.lang.reflect.InvocationTargetException;
2324
import java.time.Duration;
2425
import java.util.Arrays;
2526
import java.util.LinkedList;
@@ -170,4 +171,26 @@ public static void assertInAnnotationTransaction(boolean inTransaction) {
170171
+ " but expected in-annotation-transaction = " + inTransaction);
171172
}
172173

174+
static public Object getSecurityContext(){
175+
Object sc = null;
176+
try {
177+
Class<?> securityContextHolderClass = Class
178+
.forName("org.springframework.security.core.context.SecurityContextHolder");
179+
sc = securityContextHolderClass.getMethod("getContext").invoke(null);
180+
} catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException
181+
| InvocationTargetException cnfe) {}
182+
System.err.println(Thread.currentThread().getName() +" Util.get "+ System.identityHashCode(sc));
183+
return sc;
184+
}
185+
186+
static public void setSecurityContext(Object sc) {
187+
System.err.println(Thread.currentThread().getName() +" Util.set "+ System.identityHashCode(sc));
188+
try {
189+
Class<?> securityContextHolderClass = Class
190+
.forName("org.springframework.security.core.context.SecurityContextHolder");
191+
sc = securityContextHolderClass.getMethod("setContext", new Class[]{securityContextHolderClass}).invoke(sc);
192+
} catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException
193+
| InvocationTargetException cnfe) {}
194+
}
195+
173196
}

0 commit comments

Comments
 (0)