Merge pull request #583 from java-sec/feat-apache-httpclient-fluent

feat: 支持org.apache.httpcomponents:fluent-hc全版本增加traceId

dongtai-core/pom.xml

@@ -258,11 +258,6 @@
             <scope>test</scope>
         </dependency>
         -->
-        <dependency>
-            <groupId>org.openjdk.jmh</groupId>
-            <artifactId>jmh-generator-annprocess</artifactId>
-            <scope>test</scope>
-        </dependency>
     </dependencies>
 
 </project>

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/HttpClient.java

@@ -1,13 +1,22 @@
 package io.dongtai.iast.core.handler.hookpoint.service;
 
-import java.util.*;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
 
 public class HttpClient {
+
     private static final String JAVA_NET_URL_CONN = "sun.net.www.protocol.http.HttpURLConnection.connect()";
     private static final String JAVA_NET_URL_CONN_GET_INPUT_STREAM = "sun.net.www.protocol.http.HttpURLConnection.getInputStream()";
     private static final String JAVA_NET_URL_CONN_GET_OUTPUT_STREAM = "sun.net.www.protocol.http.HttpURLConnection.getOutputStream()";
+
+    // 草这都是哪个GA里的类啊鬼能知道啊...
     private static final String APACHE_LEGACY_HTTP_CLIENT_REQUEST_SET_URI = " org.apache.commons.httpclient.HttpMethodBase.setURI(org.apache.commons.httpclient.URI)".substring(1);
     private static final String APACHE_HTTP_CLIENT_EXECUTE = " org.apache.http.impl.client.CloseableHttpClient.doExecute(org.apache.http.HttpHost,org.apache.http.HttpRequest,org.apache.http.protocol.HttpContext)".substring(1);
+
+    // GA: org.apache.httpcomponents:fluent-hc
+    private static final String APACHE_HTTP_HTTPCOMPONENTS_EXECUTE = " org.apache.http.client.fluent.Request.execute()".substring(1);
+
     private static final String APACHE_HTTP_CLIENT5_EXECUTE = " org.apache.hc.client5.http.impl.classic.CloseableHttpClient.doExecute(org.apache.hc.core5.http.HttpHost,org.apache.hc.core5.http.ClassicHttpRequest,org.apache.hc.core5.http.protocol.HttpContext)".substring(1);
     private static final String OKHTTP_CALL_EXECUTE = "com.squareup.okhttp.Call.execute()";
     private static final String OKHTTP_CALL_ENQUEUE = "com.squareup.okhttp.Call.enqueue(com.squareup.okhttp.Callback)";
@@ -33,6 +42,7 @@ public class HttpClient {
             JAVA_NET_URL_CONN_GET_OUTPUT_STREAM,
             APACHE_LEGACY_HTTP_CLIENT_REQUEST_SET_URI,
             APACHE_HTTP_CLIENT_EXECUTE,
+            APACHE_HTTP_HTTPCOMPONENTS_EXECUTE,
             APACHE_HTTP_CLIENT5_EXECUTE,
             OKHTTP_CALL_EXECUTE,
             OKHTTP_CALL_ENQUEUE,
@@ -78,6 +88,10 @@ public static boolean matchApacheHttp4(String signature) {
     public static boolean matchApacheHttp5(String signature) {
         return APACHE_HTTP_CLIENT5_EXECUTE.equals(signature);
     }
+
+    public static boolean matchApacheHttpComponents(String signature) {
+        return APACHE_HTTP_HTTPCOMPONENTS_EXECUTE.equals(signature);
+    }
 
     public static boolean matchOkhttp(String signature) {
         return OKHTTP_SIGNATURE.contains(signature);

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/service/trace/HttpService.java

@@ -40,6 +40,8 @@ public void addTrace(MethodEvent event, PolicyNode policyNode) {
             traceId = addTraceToApacheHttpClientLegacy(event);
         } else if (HttpClient.matchOkhttp(this.matchedSignature)) {
             traceId = addTraceToOkhttp(event);
+        } else if (HttpClient.matchApacheHttpComponents(this.matchedSignature)) {
+            traceId = addTraceToApacheHttpComponents(event);
         }
 
         if (traceId != null && !traceId.isEmpty()) {
@@ -273,13 +275,66 @@ private void addHeaderToOkhttp(MethodEvent event, Map<String, String> headers) {
         }
     }
 
+    /**
+     * 添加traceId到Apache HttpComponents的请求上
+     *
+     * @param event
+     * @return
+     */
+    private String addTraceToApacheHttpComponents(MethodEvent event) {
+        Object obj = event.objectInstance;
+        if (obj == null) {
+            return null;
+        }
+        try {
+            String className = obj.getClass().getName();
+            if (!HttpClient.matchApacheHttpComponents(className)) {
+                return null;
+            }
+
+            // 关于库的版本兼容性：
+            // 在GA org.apache.httpcomponents:fluent-hc的[4.4, 4.5.14]这个区间的版本里的request字段是这个org.apache.http.client.fluent.InternalHttpRequest类型
+            // private final InternalHttpRequest request;
+            // 然后org.apache.http.client.fluent.InternalHttpRequest这个类继承的org.apache.http.message.AbstractHttpMessage上有个setHeader方法：
+            //    @Override // org.apache.http.HttpMessage
+            //    public void setHeader(String name, String value) {
+            //        Args.notNull(name, "Header name");
+            //        this.headergroup.updateHeader(new BasicHeader(name, value));
+            //    }
+            // 另外一提，org.apache.http.message.AbstractHttpMessage是在httpcomponents-httpcore:httpcore下的，它自从4.0-alpha5版本被添加了之后就没有变更过
+            //
+            // 在GA org.apache.httpcomponents:fluent-hc的[4.2, 4.4) 版本区间的request字段是org.apache.http.client.methods.HttpRequestBase类型，这个类属于依赖中的org.apache.httpcomponents:httpclient
+            // private final HttpRequestBase request;
+            // 在org.apache.httpcomponents:httpclient的[4.0.1, 4.2.6]版本区间，org.apache.http.client.methods.HttpRequestBase这个类是继承的org.apache.http.message.AbstractHttpMessage，此条分支可以与上面的合并
+            // org.apache.httpcomponents:httpclient的[4.3, 4.5.14]区间内是继承的org.apache.http.client.methods.AbstractExecutionAwareRequest
+            // org.apache.http.client.methods.AbstractExecutionAwareRequest自从4.3.4版本被添加依赖，一直继承的org.apache.http.message.AbstractHttpMessage，此条分支又可以与上面合并
+            // 所有版本的实现最终都会直接继承或者间接继承到org.apache.http.message.AbstractHttpMessage，所以下面的操作才可以统一
+
+            Field reqField = obj.getClass().getDeclaredField("request");
+            reqField.setAccessible(true);
+            Object internalHttpRequest = reqField.get(obj);
+            Method setHeaderMethod = internalHttpRequest.getClass().getMethod("setHeader", String.class, String.class);
+
+            // 然后把追踪的头加上
+            final String traceId = ContextManager.nextTraceId();
+            setHeaderMethod.invoke(internalHttpRequest, ContextManager.getHeaderKey(), traceId);
+            setHeaderMethod.invoke(internalHttpRequest, ContextManager.getParentKey(), String.valueOf(EngineManager.getAgentId()));
+            return traceId;
+        } catch (Throwable e) {
+            DongTaiLog.debug("add traceId header to apache http components failed: {}, {}", e.getMessage(), e.getCause() != null ? e.getCause().getMessage() : "");
+        }
+        return null;
+    }
+
     public static boolean validate(MethodEvent event) {
         if (HttpClient.matchJavaNetUrl(event.signature)) {
             return validateURLConnection(event);
         } else if (HttpClient.matchApacheHttp4(event.signature) || HttpClient.matchApacheHttp5(event.signature)) {
             return validateApacheHttpClient(event);
         } else if (HttpClient.matchOkhttp(event.signature)) {
             return validateOkhttp(event);
+        } else if (HttpClient.matchApacheHttpComponents(event.signature)) {
+            return validateApacheHttpComponents(event);
         }
         return true;
     }
@@ -377,4 +432,37 @@ public static boolean validateOkhttp(MethodEvent event) {
         }
         return false;
     }
+
+    /**
+     * 验证是否是合法的
+     *
+     * @param event
+     * @return
+     */
+    public static boolean validateApacheHttpComponents(MethodEvent event) {
+        Object obj = event.objectInstance;
+        if (obj == null) {
+            return false;
+        }
+        try {
+            String className = obj.getClass().getName();
+            if (!HttpClient.matchApacheHttpComponents(className)) {
+                return false;
+            }
+
+            // 关于类的版本兼容性，详见 #addTraceToApacheHttpComponents方法
+            Field reqField = obj.getClass().getDeclaredField("request");
+            reqField.setAccessible(true);
+            Object internalHttpRequest = reqField.get(obj);
+            Object header = internalHttpRequest.getClass().getMethod("getFirstHeader", String.class).invoke(internalHttpRequest, ContextManager.getHeaderKey());
+            // traceId header not exists
+            if (header == null) {
+                return true;
+            }
+        } catch (Throwable e) {
+            DongTaiLog.debug("validate apache http components failed: {}, {}", e.getMessage(), e.getCause() != null ? e.getCause().getMessage() : "");
+        }
+        return false;
+    }
+
 }