Skip to content

Commit 3e8d4ce

Browse files
committed
Add Spring Boot async context-propagation sample and e2e coverage
Demonstrate request-context propagation end to end across the JDK executor types and Spring's TaskExecutor / @async: seven /api/pets/create/async/* endpoints run the SQL insert on a worker thread and must still detect the SQLi payload. Wire the matching end2end payloads into spring_boot_postgres.
1 parent 2a53ce0 commit 3e8d4ce

4 files changed

Lines changed: 197 additions & 0 deletions

File tree

end2end/spring_boot_postgres.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,27 @@
66
safe_request=Request("/api/pets/create", body={"name": "Bobby"}),
77
unsafe_request=Request("/api/pets/create", body={"name": "Malicious Pet', 'Gru from the Minions') -- "})
88
)
9+
10+
for endpoint in [
11+
"completable-future-single",
12+
"submit-callable",
13+
"thread-pool-execute",
14+
"fork-join-submit",
15+
"scheduled-callable",
16+
"spring-task-executor",
17+
"spring-async-annotation",
18+
]:
19+
spring_boot_postgres_app.add_payload(
20+
f"sql async context propagation {endpoint}",
21+
safe_request=Request(
22+
f"/api/pets/create/async/{endpoint}",
23+
body={"name": "Bobby"}
24+
),
25+
unsafe_request=Request(
26+
f"/api/pets/create/async/{endpoint}",
27+
body={"name": "Malicious Pet', 'Gru from the Minions') -- "}
28+
)
29+
)
930
spring_boot_postgres_app.add_payload("command injection",
1031
safe_request=Request("/api/commands/execute/Johnny", method='GET'),
1132
unsafe_request=Request("/api/commands/execute/%27%3B%20sleep%202%3B%20%23%20", method='GET'),
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package com.example.demo;
2+
3+
import org.springframework.context.annotation.Bean;
4+
import org.springframework.context.annotation.Configuration;
5+
import org.springframework.scheduling.annotation.EnableAsync;
6+
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
7+
8+
import java.util.concurrent.Executor;
9+
10+
@Configuration
11+
@EnableAsync
12+
public class AsyncContextPropagationConfig {
13+
@Bean(name = "asyncContextPropagationExecutor")
14+
public Executor asyncContextPropagationExecutor() {
15+
ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
16+
executor.setThreadNamePrefix("async-context-");
17+
executor.setCorePoolSize(2);
18+
executor.setMaxPoolSize(2);
19+
executor.setQueueCapacity(10);
20+
executor.initialize();
21+
return executor;
22+
}
23+
}
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
package com.example.demo;
2+
3+
import org.springframework.beans.factory.annotation.Qualifier;
4+
import org.springframework.http.MediaType;
5+
import org.springframework.web.bind.annotation.*;
6+
7+
import java.util.concurrent.*;
8+
9+
@RestController
10+
@RequestMapping("/api/pets/create/async")
11+
public class AsyncContextPropagationController {
12+
private final Executor springExecutor;
13+
private final AsyncContextPropagationService asyncContextPropagationService;
14+
15+
public AsyncContextPropagationController(
16+
@Qualifier("asyncContextPropagationExecutor") Executor springExecutor,
17+
AsyncContextPropagationService asyncContextPropagationService
18+
) {
19+
this.springExecutor = springExecutor;
20+
this.asyncContextPropagationService = asyncContextPropagationService;
21+
}
22+
23+
private record PetCreate(String name) {}
24+
25+
@PostMapping(
26+
path = "/completable-future-single",
27+
consumes = MediaType.APPLICATION_JSON_VALUE,
28+
produces = MediaType.APPLICATION_JSON_VALUE
29+
)
30+
public PetsController.Rows completableFutureSingle(@RequestBody PetCreate pet) throws Exception {
31+
ExecutorService executor = Executors.newSingleThreadExecutor();
32+
try {
33+
return CompletableFuture
34+
.supplyAsync(() -> createPet(pet.name()), executor)
35+
.get();
36+
} finally {
37+
executor.shutdown();
38+
}
39+
}
40+
41+
@PostMapping(
42+
path = "/submit-callable",
43+
consumes = MediaType.APPLICATION_JSON_VALUE,
44+
produces = MediaType.APPLICATION_JSON_VALUE
45+
)
46+
public PetsController.Rows submitCallable(@RequestBody PetCreate pet) throws Exception {
47+
ExecutorService executor = Executors.newSingleThreadExecutor();
48+
try {
49+
return executor.submit(() -> createPet(pet.name())).get();
50+
} finally {
51+
executor.shutdown();
52+
}
53+
}
54+
55+
@PostMapping(
56+
path = "/thread-pool-execute",
57+
consumes = MediaType.APPLICATION_JSON_VALUE,
58+
produces = MediaType.APPLICATION_JSON_VALUE
59+
)
60+
public PetsController.Rows threadPoolExecute(@RequestBody PetCreate pet) throws Exception {
61+
ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1);
62+
try {
63+
CompletableFuture<PetsController.Rows> future = new CompletableFuture<>();
64+
executor.execute(() -> {
65+
try {
66+
future.complete(createPet(pet.name()));
67+
} catch (Throwable throwable) {
68+
future.completeExceptionally(throwable);
69+
}
70+
});
71+
return future.get();
72+
} finally {
73+
executor.shutdown();
74+
}
75+
}
76+
77+
@PostMapping(
78+
path = "/fork-join-submit",
79+
consumes = MediaType.APPLICATION_JSON_VALUE,
80+
produces = MediaType.APPLICATION_JSON_VALUE
81+
)
82+
public PetsController.Rows forkJoinSubmit(@RequestBody PetCreate pet) throws Exception {
83+
return ForkJoinPool.commonPool()
84+
.submit(() -> createPet(pet.name()))
85+
.get();
86+
}
87+
88+
@PostMapping(
89+
path = "/scheduled-callable",
90+
consumes = MediaType.APPLICATION_JSON_VALUE,
91+
produces = MediaType.APPLICATION_JSON_VALUE
92+
)
93+
public PetsController.Rows scheduledCallable(@RequestBody PetCreate pet) throws Exception {
94+
ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
95+
try {
96+
return executor.schedule(
97+
() -> createPet(pet.name()),
98+
1,
99+
TimeUnit.MILLISECONDS
100+
).get();
101+
} finally {
102+
executor.shutdown();
103+
}
104+
}
105+
106+
@PostMapping(
107+
path = "/spring-task-executor",
108+
consumes = MediaType.APPLICATION_JSON_VALUE,
109+
produces = MediaType.APPLICATION_JSON_VALUE
110+
)
111+
public PetsController.Rows springTaskExecutor(@RequestBody PetCreate pet) throws Exception {
112+
CompletableFuture<PetsController.Rows> future = new CompletableFuture<>();
113+
springExecutor.execute(() -> {
114+
try {
115+
future.complete(createPet(pet.name()));
116+
} catch (Throwable throwable) {
117+
future.completeExceptionally(throwable);
118+
}
119+
});
120+
return future.get();
121+
}
122+
123+
@PostMapping(
124+
path = "/spring-async-annotation",
125+
consumes = MediaType.APPLICATION_JSON_VALUE,
126+
produces = MediaType.APPLICATION_JSON_VALUE
127+
)
128+
public PetsController.Rows springAsyncAnnotation(@RequestBody PetCreate pet) throws Exception {
129+
return asyncContextPropagationService
130+
.createPetWithAsyncAnnotation(pet.name())
131+
.get();
132+
}
133+
134+
private PetsController.Rows createPet(String name) {
135+
Integer rowsCreated = DatabaseHelper.createPetByName(name);
136+
return new PetsController.Rows(rowsCreated);
137+
}
138+
}
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package com.example.demo;
2+
3+
import org.springframework.scheduling.annotation.Async;
4+
import org.springframework.stereotype.Service;
5+
6+
import java.util.concurrent.CompletableFuture;
7+
8+
@Service
9+
public class AsyncContextPropagationService {
10+
@Async("asyncContextPropagationExecutor")
11+
public CompletableFuture<PetsController.Rows> createPetWithAsyncAnnotation(String name) {
12+
Integer rowsCreated = DatabaseHelper.createPetByName(name);
13+
return CompletableFuture.completedFuture(new PetsController.Rows(rowsCreated));
14+
}
15+
}

0 commit comments

Comments
 (0)