add parent agent id header for outgoing requests

Author: lostsnow

dongtai-core/src/main/java/io/dongtai/iast/core/handler/context/ContextManager.java

@@ -22,6 +22,10 @@ public static String getHeaderKey() {
         return "dt-traceid";
     }
 
+    public static String getParentKey() {
+        return "dt-parent-agent";
+    }
+
     public static void parseTraceId(String traceId) {
         TracingContext context = TracingContext.getIncoming(traceId);
         CONTEXT.set(context);

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

@@ -40,7 +40,10 @@ public static void solveSyncInvoke(MethodEvent event, AtomicInteger invokeIdSequ
             String traceId = ContextManager.nextTraceId();
             // clear old traceId header
             addHeaderMethod.invoke(template, ContextManager.getHeaderKey(), new String[]{});
+            addHeaderMethod.invoke(template, ContextManager.getParentKey(), new String[]{});
             addHeaderMethod.invoke(template, ContextManager.getHeaderKey(), new String[]{traceId});
+            addHeaderMethod.invoke(template, ContextManager.getParentKey(),
+                    new String[]{String.valueOf(EngineManager.getAgentId())});
 
             // add to method pool
             event.source = false;

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

@@ -1,5 +1,6 @@
 package io.dongtai.iast.core.handler.hookpoint.service.trace;
 
+import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.context.ContextManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNode;
@@ -52,6 +53,8 @@ private String addTraceToJavaNetURL(MethodEvent event) {
                 final HttpURLConnection connection = (HttpURLConnection) event.objectInstance;
                 final String traceId = ContextManager.nextTraceId();
                 connection.setRequestProperty(ContextManager.getHeaderKey(), traceId);
+                connection.setRequestProperty(ContextManager.getParentKey(),
+                        String.valueOf(EngineManager.getAgentId()));
                 return traceId;
             }
         } catch (IllegalStateException ignore) {
@@ -84,6 +87,7 @@ private String addTraceToApacheHttpClient(MethodEvent event) {
             }
             final String traceId = ContextManager.nextTraceId();
             method.invoke(obj, ContextManager.getHeaderKey(), traceId);
+            method.invoke(obj, ContextManager.getParentKey(), String.valueOf(EngineManager.getAgentId()));
             return traceId;
         } catch (Throwable e) {
             DongTaiLog.debug("add traceId header to apache http client failed: {}, {}",
@@ -105,6 +109,7 @@ private String addTraceToApacheHttpClientLegacy(MethodEvent event) {
             }
             final String traceId = ContextManager.nextTraceId();
             method.invoke(obj, ContextManager.getHeaderKey(), traceId);
+            method.invoke(obj, ContextManager.getParentKey(), String.valueOf(EngineManager.getAgentId()));
             return traceId;
         } catch (Throwable e) {
             DongTaiLog.debug("add traceId header to apache legacy http client failed: {}, {}",
@@ -134,6 +139,8 @@ private String addTraceToOkhttp(MethodEvent event) {
             Method methodAddHeader = reqBuilder.getClass().getMethod("addHeader", String.class, String.class);
             final String traceId = ContextManager.nextTraceId();
             methodAddHeader.invoke(reqBuilder, ContextManager.getHeaderKey(), traceId);
+            methodAddHeader.invoke(reqBuilder, ContextManager.getParentKey(),
+                    String.valueOf(EngineManager.getAgentId()));
             Method methodBuild = reqBuilder.getClass().getMethod("build");
             Object newReq = methodBuild.invoke(reqBuilder);
             reqField.set(obj, newReq);