Skip to content

Commit 16ad60b

Browse files
willem-delbarewillem-delbare
authored
Merge pull request #107 from AikidoSec/AIK-4370
AIK-4370: Add general high level path variable wrapper
2 parents 2e810b2 + 76683ae commit 16ad60b

File tree

16 files changed

+348
-237
lines changed

16 files changed

+348
-237
lines changed

‎agent/src/main/java/dev/aikido/agent/Wrappers.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,8 @@ public final class Wrappers {
1414
private Wrappers() {}
1515
public static final List<Wrapper> WRAPPERS = Arrays.asList(
1616
new PostgresWrapper(),
17-
new SpringFrameworkWrapper(),
18-
new SpringFrameworkBodyWrapper(),
19-
new SpringFrameworkInvokeWrapper(),
17+
new SpringMVCWrapper(),
18+
new SpringControllerWrapper(),
2019
new FileWrapper(),
2120
new URLConnectionWrapper(),
2221
new InetAddressWrapper(),

‎agent/src/main/java/dev/aikido/agent/wrappers/SpringFrameworkBodyWrapper.java renamed to ‎agent/src/main/java/dev/aikido/agent/wrappers/SpringControllerWrapper.java

Lines changed: 8 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package dev.aikido.agent.wrappers;
22

3-
import dev.aikido.agent_api.collectors.RequestBodyCollector;
3+
import dev.aikido.agent_api.collectors.SpringAnnotationCollector;
44
import net.bytebuddy.asm.Advice;
55
import net.bytebuddy.description.method.MethodDescription;
66
import net.bytebuddy.description.type.TypeDescription;
@@ -15,13 +15,12 @@
1515
import static net.bytebuddy.matcher.ElementMatchers.*;
1616

1717
/* We wrap the controller functions annotated with an @RequestMapping
18-
* We check the input for @RequestBody and @RequestParam
19-
* @RequestPart currently not supported.
18+
* We check the input for @RequestBody, @RequestParam, @PathVariable, ...
2019
*/
21-
public class SpringFrameworkBodyWrapper implements Wrapper {
20+
public class SpringControllerWrapper implements Wrapper {
2221
@Override
2322
public String getName() {
24-
return SpringFrameworkBodyWrapperAdvice.class.getName();
23+
return SpringAnnotationWrapperAdvice.class.getName();
2524
}
2625

2726
@Override
@@ -43,35 +42,17 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
4342
return hasSuperType(declaresMethod(getMatcher()));
4443
}
4544

46-
private static class SpringFrameworkBodyWrapperAdvice {
45+
private static class SpringAnnotationWrapperAdvice {
4746
/* We intercept the call to the controller, arguments given to it contain user input
48-
* We check if it's annotated by @RequestBody, ... and add it as a key to our body.
47+
* We check if it's annotated by @RequestBody, @PathVariable, ... and add it as a key to our body.
4948
*/
5049
@Advice.OnMethodEnter(suppress = Throwable.class)
5150
public static void before(
5251
@Advice.Origin Executable method,
5352
@Advice.AllArguments(readOnly = false, typing = DYNAMIC) Object[] args
54-
) {
53+
) throws Exception {
5554
Parameter[] parameters = method.getParameters();
56-
for (int i = 0; i < parameters.length; i++) {
57-
Parameter parameter = parameters[i];
58-
for (Annotation annotation: parameter.getDeclaredAnnotations()) {
59-
String annotStr = annotation.toString();
60-
if (annotStr.contains("org.springframework.web.bind.annotation.RequestBody")) {
61-
// RequestBody includes all data so we report everything as one block:
62-
// Also important for API Discovery that we get the exact overview
63-
RequestBodyCollector.report(args[i]);
64-
return; // You can safely return here without missing more data
65-
}
66-
if (annotStr.contains("org.springframework.web.bind.annotation.RequestParam") ||
67-
annotStr.contains("org.springframework.web.bind.annotation.RequestPart")) {
68-
// RequestPart and RequestParam both contain partial data.
69-
String identifier = parameter.getName();
70-
RequestBodyCollector.report(identifier, args[i]);
71-
break; // You can safely exit for-loop, but we still want to scan other arguments.
72-
}
73-
}
74-
}
55+
SpringAnnotationCollector.report(parameters, args);
7556
}
7657
}
7758
}

‎agent/src/main/java/dev/aikido/agent/wrappers/SpringFrameworkInvokeWrapper.java

Lines changed: 0 additions & 53 deletions
This file was deleted.

‎agent/src/main/java/dev/aikido/agent/wrappers/SpringFrameworkWrapper.java renamed to ‎agent/src/main/java/dev/aikido/agent/wrappers/SpringMVCWrapper.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import dev.aikido.agent_api.collectors.WebRequestCollector;
44
import dev.aikido.agent_api.collectors.WebResponseCollector;
55
import dev.aikido.agent_api.context.ContextObject;
6-
import dev.aikido.agent_api.context.SpringContextObject;
6+
import dev.aikido.agent_api.context.SpringMVCContextObject;
77
import jakarta.servlet.http.Cookie;
88
import jakarta.servlet.http.HttpServletResponse;
99
import net.bytebuddy.asm.Advice;
@@ -23,16 +23,16 @@
2323
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
2424
import static net.bytebuddy.matcher.ElementMatchers.nameContains;
2525

26-
public class SpringFrameworkWrapper implements Wrapper {
27-
public static final Logger logger = LogManager.getLogger(SpringFrameworkWrapper.class);
26+
public class SpringMVCWrapper implements Wrapper {
27+
public static final Logger logger = LogManager.getLogger(SpringMVCWrapper.class);
2828

2929
@Override
3030
public String getName() {
3131
// We wrap the function doFilterInternal which gets called with
3232
// HttpServletRequest request, HttpServletResponse response
3333
// And is part of org.springframework.web.filter.RequestContextFilter
3434
// See: https://github.com/spring-projects/spring-framework/blob/4749d810db0261ce16ae5f32da6d375bb8087430/spring-web/src/main/java/org/springframework/web/filter/RequestContextFilter.java#L92
35-
return SpringFrameworkAdvice.class.getName();
35+
return SpringMVCAdvice.class.getName();
3636
}
3737
@Override
3838
public ElementMatcher<? super MethodDescription> getMatcher() {
@@ -44,7 +44,7 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
4444
return nameContains("org.springframework.web.filter.RequestContextFilter");
4545
}
4646

47-
public static class SpringFrameworkAdvice {
47+
public static class SpringMVCAdvice {
4848
// Wrapper to skip if it's inside this wrapper (i.e. our own response : )
4949
public record SkipOnWrapper(HttpServletResponse response) {};
5050
/**
@@ -76,7 +76,7 @@ public static Object interceptOnEnter(
7676
}
7777
}
7878

79-
ContextObject contextObject = new SpringContextObject(
79+
ContextObject contextObject = new SpringMVCContextObject(
8080
request.getMethod(), request.getRequestURL(), request.getRemoteAddr(),
8181
request.getParameterMap(), cookiesMap, headersMap
8282
);

‎agent_api/src/main/java/dev/aikido/agent_api/collectors/RequestBodyCollector.java

Lines changed: 0 additions & 28 deletions
This file was deleted.

‎agent_api/src/main/java/dev/aikido/agent_api/collectors/SpringAnnotationCollector.java

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
package dev.aikido.agent_api.collectors;
2+
3+
import dev.aikido.agent_api.context.Context;
4+
import dev.aikido.agent_api.context.NettyContext;
5+
import dev.aikido.agent_api.context.SpringContextObject;
6+
import dev.aikido.agent_api.context.SpringMVCContextObject;
7+
8+
import java.lang.annotation.Annotation;
9+
import java.lang.reflect.Parameter;
10+
import java.util.Map;
11+
import java.util.Optional;
12+
13+
public final class SpringAnnotationCollector {
14+
private SpringAnnotationCollector() {}
15+
private static String REQUEST_BODY = "org.springframework.web.bind.annotation.RequestBody";
16+
private static String REQUEST_PARAM = "org.springframework.web.bind.annotation.RequestParam";
17+
private static String REQUEST_PART = "org.springframework.web.bind.annotation.RequestPart";
18+
private static String PATH_VARIABLE = "org.springframework.web.bind.annotation.PathVariable";
19+
20+
/** report(...)
21+
* Handles Springs parameters, can include @PathVariable, @RequestBody, ...
22+
*/
23+
public static void report(Parameter[] parameters, Object[] values) throws Exception {
24+
if (parameters.length != values.length) {
25+
// This exception gets caught by Byte Buddy.
26+
throw new Exception("Length of parameters and values should match!");
27+
}
28+
for (int i = 0; i < parameters.length; i++) {
29+
report(parameters[i], values[i]);
30+
}
31+
}
32+
33+
public static void report(Parameter parameter, Object value) {
34+
SpringContextObject context = (SpringContextObject) Context.get();
35+
if (context == null) {
36+
return;
37+
}
38+
39+
for (Annotation annotation: parameter.getDeclaredAnnotations()) {
40+
String annotStr = annotation.annotationType().getName();
41+
if (annotStr.contains(REQUEST_BODY)) {
42+
// RequestBody includes all data so we report everything as one block:
43+
// Also important for API Discovery that we get the exact overview
44+
context.setBody(value);
45+
break;
46+
} else if (annotStr.equals(REQUEST_PARAM) || annotStr.equals(REQUEST_PART)) {
47+
// RequestPart and RequestParam both contain partial data.
48+
String identifier = parameter.getName();
49+
context.setBodyElement(identifier, value);
50+
break;
51+
} else if(annotStr.equals(PATH_VARIABLE)) {
52+
String identifier = parameter.getName();
53+
if (value instanceof Map<?, ?> paramsMap) {
54+
for (Map.Entry<?, ?> entry: paramsMap.entrySet()) {
55+
if (entry.getKey() instanceof String key && entry.getValue() instanceof String valueStr) {
56+
context.setParameter(key, valueStr);
57+
}
58+
}
59+
} else if (value instanceof String valueStr) {
60+
context.setParameter(identifier, valueStr);
61+
} else if (value instanceof Optional<?> valueOpt) {
62+
if(valueOpt.isPresent()) {
63+
if (valueOpt.get() instanceof String valueStr) {
64+
context.setParameter(identifier, valueStr);
65+
}
66+
}
67+
}
68+
break;
69+
}
70+
}
71+
Context.set(context); // Store context.
72+
}
73+
}

‎agent_api/src/main/java/dev/aikido/agent_api/context/NettyContext.java

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import static dev.aikido.agent_api.helpers.net.ProxyForwardedParser.getIpFromRequest;
1010
import static dev.aikido.agent_api.helpers.url.BuildRouteFromUrl.buildRouteFromUrl;
1111

12-
public class NettyContext extends ContextObject {
12+
public class NettyContext extends SpringContextObject {
1313
public NettyContext(
1414
String method, String uri, InetSocketAddress rawIp,
1515
HashMap<String, List<String>> cookies,
@@ -27,9 +27,6 @@ public NettyContext(
2727
this.remoteAddress = getIpFromRequest(rawIp.getAddress().getHostAddress(), this.headers);
2828
this.source = "ReactorNetty";
2929
this.redirectStartNodes = new ArrayList<>();
30-
31-
// We don't have access yet to the route parameters.
32-
this.params = null;
3330
}
3431

3532
private static HashMap<String, String> extractHeaders(List<Map.Entry<String, String>> entries) {

‎agent_api/src/main/java/dev/aikido/agent_api/context/SpringContextObject.java

Lines changed: 14 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,58 +1,29 @@
11
package dev.aikido.agent_api.context;
22

3-
import java.util.*;
4-
5-
import static dev.aikido.agent_api.helpers.net.ProxyForwardedParser.getIpFromRequest;
6-
import static dev.aikido.agent_api.helpers.url.BuildRouteFromUrl.buildRouteFromUrl;
3+
import java.util.HashMap;
4+
import java.util.Map;
75

86
public class SpringContextObject extends ContextObject {
97
// We use this map for when @RequestBody does not get used :
108
protected transient Map<String, Object> bodyMap = new HashMap<>();
11-
public SpringContextObject(
12-
String method, StringBuffer url, String rawIp, Map<String, String[]> queryParams,
13-
HashMap<String, List<String>> cookies, HashMap<String, String> headers
14-
) {
15-
this.method = method;
16-
if (url != null) {
17-
this.url = url.toString();
18-
}
19-
this.query = extractQueryParameters(queryParams);
20-
this.cookies = cookies;
21-
this.headers = headers;
22-
this.route = buildRouteFromUrl(this.url);
23-
this.remoteAddress = getIpFromRequest(rawIp, this.headers);
24-
this.source = "SpringFramework";
25-
this.redirectStartNodes = new ArrayList<>();
269

27-
// We don't have access yet to the route parameters: doFilter() is called before the Controller
28-
// So the parameters will be set later.
29-
this.params = null;
10+
// We don't have access yet to the route parameters: doFilter() is called before the Controller
11+
// So the parameters will be set later.
12+
protected transient Map<String, String> params = new HashMap<>();
13+
14+
public void setParameter(String key, String value) {
15+
this.params.put(key, value);
16+
this.cache.remove("routeParams"); // Reset cache
3017
}
18+
3119
@Override
32-
public Object getBody() {
33-
if (this.body != null) {
34-
// @RequestBody was used, all data is available :
35-
return this.body;
36-
}
37-
return this.bodyMap; // Use the selected fields that were extracted.
20+
public Map<String, String> getParams() {
21+
return params;
3822
}
23+
3924
public void setBodyElement(String key, Object value) {
4025
bodyMap.put(key, value);
4126
cache.remove("body"); // Reset body cache.
4227
}
43-
public void setParams(Object params) {
44-
this.params = params;
45-
this.cache.remove("routeParams"); // Reset cache
46-
}
47-
48-
private static HashMap<String, List<String>> extractQueryParameters(Map<String, String[]> parameterMap) {
49-
HashMap<String, List<String>> query = new HashMap<>();
50-
51-
for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
52-
// Convert String[] to List<String>
53-
List<String> list = Arrays.asList(entry.getValue());
54-
query.put(entry.getKey(), list);
55-
}
56-
return query;
57-
}
5828
}
29+

0 commit comments

Comments
 (0)