Skip to content

Commit b96d425

Browse files
bitterpanda63bitterpanda63
committed
Update wrappers 1/2
1 parent fc86a1c commit b96d425

9 files changed

Lines changed: 37 additions & 318 deletions

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

Lines changed: 4 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,16 @@
11
package dev.aikido.agent.wrappers;
22

3+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
34
import net.bytebuddy.asm.Advice;
45
import net.bytebuddy.description.method.MethodDescription;
56
import net.bytebuddy.description.type.TypeDescription;
67
import net.bytebuddy.matcher.ElementMatcher;
78
import net.bytebuddy.matcher.ElementMatchers;
89

9-
import java.lang.reflect.Method;
10-
import java.net.HttpURLConnection;
11-
import java.net.MalformedURLException;
1210
import java.net.URL;
13-
import java.net.URLClassLoader;
1411
import java.net.http.HttpClient;
1512
import java.net.http.HttpRequest;
1613

17-
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
18-
1914
public class HttpClientSendWrapper implements Wrapper {
2015
public String getName() {
2116
// Wrap send(HttpRequest req, ...) and sendAsync(HttpRequest req, ...) on HttpClient instance
@@ -32,37 +27,15 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
3227
return ElementMatchers.isSubTypeOf(HttpClient.class);
3328
}
3429
public static class SendFunctionAdvice {
35-
// Since we have to wrap a native Java Class stuff gets more complicated
36-
// The classpath is not the same anymore, and we can't import our modules directly.
37-
// To bypass this issue we load collectors from a .jar file
3830
@Advice.OnMethodEnter(suppress = Throwable.class)
3931
public static void before(
4032
@Advice.Argument(0) HttpRequest httpRequest
41-
) throws Exception {
33+
) throws Throwable {
4234
if (httpRequest == null || httpRequest.uri() == null) {
4335
return;
4436
}
45-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
46-
URLClassLoader classLoader = null;
47-
try {
48-
URL[] urls = { new URL(jarFilePath) };
49-
classLoader = new URLClassLoader(urls);
50-
} catch (MalformedURLException ignored) {}
51-
if (classLoader == null) {
52-
return;
53-
}
54-
55-
// Load the class from the JAR
56-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
57-
58-
// Run report with "argument"
59-
for (Method method2: clazz.getMethods()) {
60-
if(method2.getName().equals("report")) {
61-
method2.invoke(null, httpRequest.uri().toURL());
62-
break;
63-
}
64-
}
65-
classLoader.close(); // Close the class loader
37+
URL url = httpRequest.uri().toURL();
38+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.URL_COLLECTOR_REPORT, url);
6639
}
6740
}
6841
}

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

Lines changed: 5 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,16 @@
11
package dev.aikido.agent.wrappers;
22

3+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
34
import net.bytebuddy.asm.Advice;
45
import net.bytebuddy.description.method.MethodDescription;
56
import net.bytebuddy.description.type.TypeDescription;
67
import net.bytebuddy.matcher.ElementMatcher;
78
import net.bytebuddy.matcher.ElementMatchers;
89

9-
import java.lang.reflect.Method;
1010
import java.net.*;
1111
import java.net.http.HttpClient;
1212
import java.net.http.HttpRequest;
1313

14-
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
15-
1614
public class HttpClientWrapper implements Wrapper {
1715
public String getName() {
1816
// Wrap getResponseCode function which executes HTTP requests
@@ -28,35 +26,14 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
2826
return ElementMatchers.isSubTypeOf(HttpClient.class);
2927
}
3028
public static class HttpClientAdvice {
31-
// Since we have to wrap a native Java Class stuff gets more complicated
32-
// The classpath is not the same anymore, and we can't import our modules directly.
33-
// To bypass this issue we load collectors from a .jar file, specified with the AIKIDO_DIRECTORY env variable
3429
@Advice.OnMethodEnter(suppress = Throwable.class)
3530
public static void before(
36-
@Advice.This(typing = DYNAMIC, optional = true) Object target,
3731
@Advice.Argument(0) Object uriObject,
3832
@Advice.Argument(2) Object httpRequestObject
39-
) throws Exception {
40-
URI uri = (URI) uriObject;
41-
HttpRequest httpRequest = (HttpRequest) httpRequestObject;
42-
43-
// Call to collector :
44-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
45-
URL[] urls = { new URL(jarFilePath) };
46-
URLClassLoader classLoader = new URLClassLoader(urls);
47-
48-
// Load the class from the JAR
49-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.RedirectCollector");
50-
51-
// Run report func with its arguments
52-
for (Method method2: clazz.getMethods()) {
53-
if(method2.getName().equals("report")) {
54-
URL originUrl = httpRequest.uri().toURL();
55-
method2.invoke(null, originUrl, uri.toURL());
56-
break;
57-
}
58-
}
59-
classLoader.close(); // Close the class loader
33+
) throws Throwable {
34+
URL origin = ((HttpRequest) httpRequestObject).uri().toURL();
35+
URL dest = ((URI) uriObject).toURL();
36+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.REDIRECT_COLLECTOR_REPORT, origin, dest);
6037
}
6138
}
6239
}

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

Lines changed: 3 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,13 @@
11
package dev.aikido.agent.wrappers;
22

3+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
34
import net.bytebuddy.asm.Advice;
45
import net.bytebuddy.description.method.MethodDescription;
56
import net.bytebuddy.description.type.TypeDescription;
67
import net.bytebuddy.matcher.ElementMatcher;
78
import net.bytebuddy.matcher.ElementMatchers;
89

9-
import java.lang.reflect.InvocationTargetException;
10-
import java.lang.reflect.Method;
1110
import java.net.*;
12-
import java.net.http.HttpClient;
13-
import java.util.List;
14-
import java.util.Map;
1511

1612
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
1713

@@ -32,38 +28,16 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
3228
return ElementMatchers.isSubTypeOf(HttpURLConnection.class);
3329
}
3430
public static class FollowRedirect0Advice {
35-
// Since we have to wrap a native Java Class stuff gets more complicated
36-
// The classpath is not the same anymore, and we can't import our modules directly.
37-
// To bypass this issue we load collectors from a .jar file, specified with the AIKIDO_DIRECTORY env variable
3831
@Advice.OnMethodEnter(suppress = Throwable.class)
3932
public static void before(
4033
@Advice.This(typing = DYNAMIC, optional = true) HttpURLConnection target,
4134
@Advice.Argument(2) URL destUrl
42-
) throws Exception {
35+
) throws Throwable {
4336
URL origin = target.getURL();
4437
if (origin == null || destUrl == null) {
4538
return;
4639
}
47-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
48-
URLClassLoader classLoader = null;
49-
try {
50-
URL[] urls = { new URL(jarFilePath) };
51-
classLoader = new URLClassLoader(urls);
52-
} catch (MalformedURLException ignored) {}
53-
if (classLoader == null) {
54-
return;
55-
}
56-
// Load the class from the JAR
57-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.RedirectCollector");
58-
59-
// Run report with "argument"
60-
for (Method method2: clazz.getMethods()) {
61-
if(method2.getName().equals("report")) {
62-
method2.invoke(null, origin, destUrl);
63-
break;
64-
}
65-
}
66-
classLoader.close(); // Close the class loader
40+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.REDIRECT_COLLECTOR_REPORT, origin, destUrl);
6741
}
6842
}
6943
}

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

Lines changed: 7 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
package dev.aikido.agent.wrappers;
2+
3+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
24
import net.bytebuddy.asm.Advice;
35
import net.bytebuddy.description.method.MethodDescription;
46
import net.bytebuddy.description.type.TypeDescription;
57
import net.bytebuddy.matcher.ElementMatcher;
6-
import net.bytebuddy.matcher.ElementMatchers;
78

8-
import java.lang.reflect.Method;
99
import java.net.*;
1010

1111
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
@@ -18,46 +18,20 @@ public String getName() {
1818
return ConstructorAdvice.class.getName();
1919
}
2020
public ElementMatcher<? super MethodDescription> getMatcher() {
21-
return isConstructor().and(isDeclaredBy(is(HttpURLConnection.class)));
21+
return isDeclaredBy(URL.class).and(named("openConnection"));
2222
}
2323

2424
@Override
2525
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
26-
return ElementMatchers.is(HttpURLConnection.class);
26+
return is(URL.class);
2727
}
2828

2929
public static class ConstructorAdvice {
30-
// Since we have to wrap a native Java Class stuff gets more complicated
31-
// The classpath is not the same anymore, and we can't import our modules directly.
32-
// To bypass this issue we load collectors from a .jar file
3330
@Advice.OnMethodExit(suppress = Throwable.class)
3431
public static void before(
35-
@Advice.This(typing = DYNAMIC) Object target
36-
) throws Exception {
37-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
38-
URLClassLoader classLoader = null;
39-
try {
40-
URL[] urls = { new URL(jarFilePath) };
41-
classLoader = new URLClassLoader(urls);
42-
} catch (MalformedURLException ignored) {}
43-
if (classLoader == null) {
44-
return;
45-
}
46-
URL url = ((HttpURLConnection) target).getURL();
47-
if (target == null || url == null) {
48-
return;
49-
}
50-
// Load the class from the JAR
51-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.URLCollector");
52-
53-
// Run report with "argument"
54-
for (Method method2: clazz.getMethods()) {
55-
if(method2.getName().equals("report")) {
56-
method2.invoke(null, url);
57-
break;
58-
}
59-
}
60-
classLoader.close(); // Close the class loader
32+
@Advice.This(typing = DYNAMIC) URL url
33+
) throws Throwable {
34+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.URL_COLLECTOR_REPORT, url);
6135
}
6236
}
6337
}

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

Lines changed: 2 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,12 @@
11
package dev.aikido.agent.wrappers;
2+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
23
import net.bytebuddy.asm.Advice;
34
import net.bytebuddy.description.method.MethodDescription;
45
import net.bytebuddy.description.type.TypeDescription;
56
import net.bytebuddy.matcher.ElementMatcher;
67
import net.bytebuddy.matcher.ElementMatchers;
78

8-
import java.lang.reflect.Executable;
9-
import java.lang.reflect.InvocationTargetException;
10-
import java.lang.reflect.Method;
119
import java.net.InetAddress;
12-
import java.net.MalformedURLException;
13-
import java.net.URL;
14-
import java.net.URLClassLoader;
15-
16-
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
1710

1811
public class InetAddressWrapper implements Wrapper {
1912
public String getName() {
@@ -29,44 +22,12 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
2922
return ElementMatchers.isSubTypeOf(InetAddress.class);
3023
}
3124
public static class InetAdvice {
32-
// Since we have to wrap a native Java Class stuff gets more complicated
33-
// The classpath is not the same anymore, and we can't import our modules directly.
34-
// To bypass this issue we load collectors from a .jar file
3525
@Advice.OnMethodExit
3626
public static void after(
3727
@Advice.Argument(0) String hostname,
3828
@Advice.Return InetAddress[] inetAddresses
3929
) throws Throwable {
40-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
41-
URLClassLoader classLoader = null;
42-
try {
43-
URL[] urls = { new URL(jarFilePath) };
44-
classLoader = new URLClassLoader(urls);
45-
} catch (MalformedURLException ignored) {}
46-
if (classLoader == null) {
47-
return;
48-
}
49-
50-
try {
51-
// Load the class from the JAR
52-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.DNSRecordCollector");
53-
54-
// Run report with "argument"
55-
for (Method method2: clazz.getMethods()) {
56-
if(method2.getName().equals("report")) {
57-
method2.invoke(null, hostname, inetAddresses);
58-
break;
59-
}
60-
}
61-
classLoader.close(); // Close the class loader
62-
} catch (InvocationTargetException invocationTargetException) {
63-
if(invocationTargetException.getCause().toString().startsWith("dev.aikido.agent_api.vulnerabilities")) {
64-
throw invocationTargetException.getCause();
65-
}
66-
// Ignore non-aikido throwables.
67-
} catch(Throwable e) {
68-
System.out.println("AIKIDO: " + e.getMessage());
69-
}
30+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.DNS_RECORD_COLLECTOR_REPORT, hostname, inetAddresses);
7031
}
7132
}
7233
}

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

Lines changed: 3 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,14 @@
11
package dev.aikido.agent.wrappers;
2+
import dev.aikido.agent_bootstrap.AikidoBootstrapClass;
23
import net.bytebuddy.asm.Advice;
34
import net.bytebuddy.description.method.MethodDescription;
45
import net.bytebuddy.description.type.TypeDescription;
56
import net.bytebuddy.matcher.ElementMatcher;
67
import net.bytebuddy.matcher.ElementMatchers;
78

89
import java.lang.reflect.Executable;
9-
import java.lang.reflect.InvocationTargetException;
10-
import java.lang.reflect.Method;
11-
import java.net.MalformedURLException;
12-
import java.net.URL;
13-
import java.net.URLClassLoader;
1410
import java.nio.file.Path;
1511

16-
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
17-
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.of;
1812
import static net.bytebuddy.matcher.ElementMatchers.*;
1913

2014
/**
@@ -40,41 +34,13 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
4034
}
4135

4236
public static class PathAdvice {
43-
// Since we have to wrap a native Java Class stuff gets more complicated
44-
// The classpath is not the same anymore, and we can't import our modules directly.
45-
// To bypass this issue we load collectors from a .jar file
4637
@Advice.OnMethodEnter
4738
public static void before(
4839
@Advice.Origin Executable method,
4940
@Advice.Argument(value = 0, optional = true) Object argument
5041
) throws Throwable {
51-
String jarFilePath = System.getProperty("AIK_agent_api_jar");
52-
URLClassLoader classLoader = null;
53-
try {
54-
URL[] urls = { new URL(jarFilePath) };
55-
classLoader = new URLClassLoader(urls);
56-
} catch (MalformedURLException ignored) {}
57-
if (classLoader == null) {
58-
return;
59-
}
60-
61-
try {
62-
// Load the class from the JAR
63-
Class<?> clazz = classLoader.loadClass("dev.aikido.agent_api.collectors.FileCollector");
64-
65-
// Run report with "argument"
66-
Method reportMethod = clazz.getMethod("report", Object.class, String.class);
67-
String op = "java.nio.file.Path." + method.getName();
68-
reportMethod.invoke(null, argument, op);
69-
} catch (InvocationTargetException invocationTargetException) {
70-
if(invocationTargetException.getCause().toString().startsWith("dev.aikido.agent_api.vulnerabilities")) {
71-
throw invocationTargetException.getCause();
72-
}
73-
// Ignore non-aikido throwables.
74-
} catch(Throwable e) {
75-
System.out.println("AIKIDO: " + e.getMessage());
76-
}
77-
classLoader.close(); // Close the class loader
42+
String op = "java.nio.file.Path." + method.getName();
43+
AikidoBootstrapClass.invoke(AikidoBootstrapClass.FILE_COLLECTOR_REPORT, argument, op);
7844
}
7945
}
8046
}

0 commit comments

Comments
 (0)