Skip to content

Commit 7caa870

Browse files
bitterpanda63bitterpanda63
authored
Merge pull request #283 from AikidoSec/dangerous-bodies-update
draft: add DangerousBodyException
2 parents 31e9c70 + b595c3e commit 7caa870

6 files changed

Lines changed: 127 additions & 18 deletions

File tree

‎agent_api/src/main/java/dev/aikido/agent_api/helpers/extraction/StringExtractor.java‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package dev.aikido.agent_api.helpers.extraction;
22

33
import dev.aikido.agent_api.helpers.patterns.LooksLikeJWT;
4+
import dev.aikido.agent_api.vulnerabilities.DangerousBodyException;
45
import java.lang.reflect.Field;
56
import java.lang.reflect.Modifier;
67
import java.util.*;
@@ -9,31 +10,35 @@
910
import static dev.aikido.agent_api.helpers.patterns.PrimitiveType.isPrimitiveType;
1011

1112
public class StringExtractor {
13+
private static final int MAX_DEPTH = 1024;
1214
// Ensures that we don't get recursion :
1315
Set<Object> scanned = new HashSet<>();
1416
public static Map<String, String> extractStringsFromObject(Object obj) {
15-
return new StringExtractor().extractStringsRecursive(obj, new ArrayList<>());
17+
return new StringExtractor().extractStringsRecursive(obj, new ArrayList<>(), 0);
1618
}
17-
private Map<String, String> extractStringsRecursive(Object target, ArrayList<PathBuilder.PathPart> pathToPayload) {
19+
private Map<String, String> extractStringsRecursive(Object target, ArrayList<PathBuilder.PathPart> pathToPayload, int depth) {
20+
if (depth > MAX_DEPTH) {
21+
throw DangerousBodyException.bodyTooDeep();
22+
}
1823
HashMap<String, String> result = new HashMap<>();
1924
if (target == null || scanned.contains(target)) {
2025
return Map.of(); // Do not rescan objects, because this might lead to recursion.
2126
}
2227
scanned.add(target);
2328

2429
if (target instanceof String targetString) {
25-
result.putAll(extractStringsFromString(targetString, pathToPayload));
30+
result.putAll(extractStringsFromString(targetString, pathToPayload, depth));
2631
} else if (target instanceof Collection<?> || target.getClass().isArray()) {
27-
result.putAll(extractStringsFromArray(target, pathToPayload));
32+
result.putAll(extractStringsFromArray(target, pathToPayload, depth));
2833
} else if (target instanceof Map<?, ?> targetMap) {
29-
result.putAll(extractStringsFromMap(targetMap, pathToPayload));
34+
result.putAll(extractStringsFromMap(targetMap, pathToPayload, depth));
3035
} else if (!isPrimitiveType(target)) { // Stop algorithm if it's a primitive type.
31-
result.putAll(extractStringsFromStructure(target, pathToPayload));
36+
result.putAll(extractStringsFromStructure(target, pathToPayload, depth));
3237
}
3338
return result;
3439
}
3540

36-
private Map<String, String> extractStringsFromString(String target, ArrayList<PathBuilder.PathPart> pathToPayload) {
41+
private Map<String, String> extractStringsFromString(String target, ArrayList<PathBuilder.PathPart> pathToPayload, int depth) {
3742
HashMap<String, String> result = new HashMap<>();
3843
result.put(target, buildPathToPayload(pathToPayload));
3944

@@ -42,7 +47,7 @@ private Map<String, String> extractStringsFromString(String target, ArrayList<Pa
4247
if (jwtResult.success()) {
4348
ArrayList<PathBuilder.PathPart> newPathToPayload = new ArrayList<>(pathToPayload);
4449
newPathToPayload.add(new PathBuilder.PathPart("jwt"));
45-
Map<String, String> resultsFromJWT = extractStringsRecursive(jwtResult.payload(), newPathToPayload);
50+
Map<String, String> resultsFromJWT = extractStringsRecursive(jwtResult.payload(), newPathToPayload, depth + 1);
4651
for (Map.Entry<String, String> entry : resultsFromJWT.entrySet()) {
4752
String key = entry.getKey();
4853
String value = entry.getValue();
@@ -57,27 +62,27 @@ private Map<String, String> extractStringsFromString(String target, ArrayList<Pa
5762
return result;
5863
}
5964

60-
private Map<String, String> extractStringsFromArray(Object target, ArrayList<PathBuilder.PathPart> pathToPayload) {
65+
private Map<String, String> extractStringsFromArray(Object target, ArrayList<PathBuilder.PathPart> pathToPayload, int depth) {
6166
HashMap<String, String> result = new HashMap<>();
6267
if (target instanceof Collection<?> targetCollection) {
6368
int index = 0;
6469
for (Object element : (Collection<?>) targetCollection) {
6570
ArrayList<PathBuilder.PathPart> newPathToPayload = new ArrayList<>(pathToPayload);
6671
newPathToPayload.add(new PathBuilder.PathPart("array", index));
67-
result.putAll(extractStringsRecursive(element, newPathToPayload));
72+
result.putAll(extractStringsRecursive(element, newPathToPayload, depth + 1));
6873
index++;
6974
}
7075
} else if (target instanceof Object[] targetArray) {
7176
for (int i = 0; i < targetArray.length; i++) {
7277
ArrayList<PathBuilder.PathPart> newPathToPayload = new ArrayList<>(pathToPayload);
7378
newPathToPayload.add(new PathBuilder.PathPart("array", i));
74-
result.putAll(extractStringsRecursive(targetArray[i], newPathToPayload));
79+
result.putAll(extractStringsRecursive(targetArray[i], newPathToPayload, depth + 1));
7580
}
7681
}
7782
return result;
7883
}
7984

80-
private Map<String, String> extractStringsFromMap(Map<?, ?> target, ArrayList<PathBuilder.PathPart> pathToPayload) {
85+
private Map<String, String> extractStringsFromMap(Map<?, ?> target, ArrayList<PathBuilder.PathPart> pathToPayload, int depth) {
8186
HashMap<String, String> result = new HashMap<>();
8287
for (Object key : target.keySet()) {
8388
if (key instanceof String stringKey) {
@@ -89,12 +94,12 @@ private Map<String, String> extractStringsFromMap(Map<?, ?> target, ArrayList<Pa
8994
} else {
9095
newPathToPayload.add(new PathBuilder.PathPart("object", key.toString()));
9196
}
92-
result.putAll(extractStringsRecursive(target.get(key), newPathToPayload));
97+
result.putAll(extractStringsRecursive(target.get(key), newPathToPayload, depth + 1));
9398
}
9499
return result;
95100
}
96101

97-
private Map<String, String> extractStringsFromStructure(Object target, ArrayList<PathBuilder.PathPart> pathToPayload) {
102+
private Map<String, String> extractStringsFromStructure(Object target, ArrayList<PathBuilder.PathPart> pathToPayload, int depth) {
98103
HashMap<String, String> result = new HashMap<>();
99104
Field[] fields = target.getClass().getDeclaredFields();
100105
for (Field field : fields) {
@@ -106,7 +111,9 @@ private Map<String, String> extractStringsFromStructure(Object target, ArrayList
106111
Object fieldValue = field.get(target);
107112
ArrayList<PathBuilder.PathPart> newPathToPayload = new ArrayList<>(pathToPayload);
108113
newPathToPayload.add(new PathBuilder.PathPart("object", field.getName()));
109-
result.putAll(extractStringsRecursive(fieldValue, newPathToPayload));
114+
result.putAll(extractStringsRecursive(fieldValue, newPathToPayload, depth + 1));
115+
} catch (DangerousBodyException e) {
116+
throw e;
110117
} catch (IllegalAccessException | RuntimeException e) {
111118
// pass-through
112119
}

‎agent_api/src/main/java/dev/aikido/agent_api/helpers/patterns/LooksLikeJWT.java‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,16 @@
44
import com.google.gson.Gson;
55
import com.google.gson.reflect.TypeToken;
66

7+
import dev.aikido.agent_api.vulnerabilities.DangerousBodyException;
8+
79
import java.util.Map;
810
import java.util.Objects;
911

1012
public final class LooksLikeJWT {
1113
private LooksLikeJWT() {}
1214

15+
public static final int MAX_JWT_PAYLOAD_BYTES = 8 * 1024;
16+
1317
public static Result tryDecodeAsJwt(String jwt) {
1418
// Check if the JWT contains the required parts
1519
if (jwt == null || !jwt.contains(".")) {
@@ -23,14 +27,24 @@ public static Result tryDecodeAsJwt(String jwt) {
2327
return new Result(false, null);
2428
}
2529

30+
byte[] decoded;
2631
try {
27-
// Decode the middle part (payload) of the JWT
28-
String payload = new String(Base64.getUrlDecoder().decode(parts[1]));
32+
decoded = Base64.getUrlDecoder().decode(parts[1]);
33+
} catch (IllegalArgumentException ignored) {
34+
return new Result(false, null);
35+
}
2936

37+
if (decoded.length > MAX_JWT_PAYLOAD_BYTES) {
38+
throw DangerousBodyException.jwtTooLarge();
39+
}
40+
41+
String payload = new String(decoded);
42+
try {
3043
Gson gson = new Gson();
3144
Map<String, Object> jwtPayload = gson.fromJson(payload, new TypeToken<Map<String, Object>>(){}.getType());
32-
3345
return new Result(true, jwtPayload);
46+
} catch (StackOverflowError soe) {
47+
throw DangerousBodyException.jwtTooLarge();
3448
} catch (Exception ignored) {
3549
return new Result(false, null);
3650
}

‎agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/DangerousBodyException.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package dev.aikido.agent_api.vulnerabilities;
2+
3+
public class DangerousBodyException extends AikidoException {
4+
public DangerousBodyException(String reason) {
5+
super(generateDefaultMessage("Dangerous Body") + ": " + reason);
6+
}
7+
8+
public static DangerousBodyException jwtTooLarge() {
9+
return new DangerousBodyException("JWT payload too large");
10+
}
11+
12+
public static DangerousBodyException bodyTooDeep() {
13+
return new DangerousBodyException("Body is too deeply nested to scan");
14+
}
15+
}

‎agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/Scanner.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ public static void scanForGivenVulnerability(Vulnerabilities.Vulnerability vulne
5050
break;
5151
}
5252
}
53+
} catch (AikidoException ae) {
54+
exception = Optional.of(ae);
5355
} catch (Throwable e) {
5456
logger.debug(e);
5557
}

‎agent_api/src/test/java/helpers/LooksLikeJWTTest.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
package helpers;
22

33
import dev.aikido.agent_api.helpers.patterns.LooksLikeJWT;
4+
import dev.aikido.agent_api.vulnerabilities.DangerousBodyException;
45
import org.junit.jupiter.api.Test;
56

7+
import java.util.Base64;
68
import java.util.HashMap;
79
import java.util.Map;
810

911
import static org.junit.jupiter.api.Assertions.assertEquals;
12+
import static org.junit.jupiter.api.Assertions.assertThrows;
1013

1114
public class LooksLikeJWTTest {
1215

@@ -52,4 +55,35 @@ void testReturnsDecodedJwtForValidJwtWithBearerPrefix() {
5255
expectedPayload.put("iat", 1.516239022E9);
5356
assertEquals(new LooksLikeJWT.Result(true, expectedPayload), LooksLikeJWT.tryDecodeAsJwt(validJwtWithBearer));
5457
}
58+
59+
private static String buildJwt(String payloadJson) {
60+
String payloadB64 = Base64.getUrlEncoder().withoutPadding()
61+
.encodeToString(payloadJson.getBytes());
62+
return "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9." + payloadB64 + ".sig";
63+
}
64+
65+
@Test
66+
void testThrowsDangerousBodyExceptionForOversizedPayload() {
67+
StringBuilder sb = new StringBuilder("{\"k\":\"");
68+
for (int i = 0; i < LooksLikeJWT.MAX_JWT_PAYLOAD_BYTES + 10; i++) {
69+
sb.append('a');
70+
}
71+
sb.append("\"}");
72+
String jwt = buildJwt(sb.toString());
73+
assertThrows(DangerousBodyException.class, () -> LooksLikeJWT.tryDecodeAsJwt(jwt));
74+
}
75+
76+
@Test
77+
void testThrowsDangerousBodyExceptionForDeeplyNestedPayload() {
78+
int depth = 7000;
79+
StringBuilder open = new StringBuilder();
80+
StringBuilder close = new StringBuilder();
81+
for (int i = 0; i < depth; i++) {
82+
open.append("{\"a\":");
83+
close.append("}");
84+
}
85+
String payload = open.toString() + "1" + close.toString();
86+
String jwt = buildJwt(payload);
87+
assertThrows(DangerousBodyException.class, () -> LooksLikeJWT.tryDecodeAsJwt(jwt));
88+
}
5589
}

‎agent_api/src/test/java/helpers/StringExtractorTest.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import dev.aikido.agent_api.api_discovery.DataSchemaGenerator;
55
import dev.aikido.agent_api.api_discovery.DataSchemaItem;
66
import dev.aikido.agent_api.api_discovery.DataSchemaType;
7+
import dev.aikido.agent_api.vulnerabilities.DangerousBodyException;
78
import org.junit.jupiter.api.Test;
89

910
import static dev.aikido.agent_api.helpers.extraction.StringExtractor.extractStringsFromObject;
@@ -407,6 +408,42 @@ public void testItChecksScannedClasses() {
407408
assertEquals(".important_record.a", result.get("Hello World"));
408409
}
409410

411+
@Test
412+
public void testThrowsDangerousBodyExceptionForTooDeepNesting() {
413+
Map<String, Object> root = new HashMap<>();
414+
Map<String, Object> current = root;
415+
for (int i = 0; i < 1100; i++) {
416+
Map<String, Object> next = new HashMap<>();
417+
current.put("a", next);
418+
current = next;
419+
}
420+
current.put("leaf", "x");
421+
assertThrows(DangerousBodyException.class, () -> extractStringsFromObject(root));
422+
}
423+
424+
@Test
425+
public void testDoesNotThrowForModestNesting() {
426+
Map<String, Object> root = new HashMap<>();
427+
Map<String, Object> current = root;
428+
for (int i = 0; i < 500; i++) {
429+
Map<String, Object> next = new HashMap<>();
430+
current.put("a", next);
431+
current = next;
432+
}
433+
current.put("leaf", "x");
434+
Map<String, String> result = extractStringsFromObject(root);
435+
assertEquals("x", findKeyEndingWithLeafValue(result));
436+
}
437+
438+
private static String findKeyEndingWithLeafValue(Map<String, String> result) {
439+
for (Map.Entry<String, String> e : result.entrySet()) {
440+
if ("x".equals(e.getKey())) {
441+
return e.getKey();
442+
}
443+
}
444+
return null;
445+
}
446+
410447
@Test
411448
public void testItChecksScannedObjects() {
412449
Map<String, Object> input = new HashMap<>();

0 commit comments

Comments
 (0)