Merge pull request #594 from 15911075183ma/fix-feign-bug

Fix feign bug

Author: mazp99

dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java

@@ -1,9 +1,13 @@
 package io.dongtai.iast.agent.manager;
 
-import io.dongtai.iast.agent.*;
+import io.dongtai.iast.agent.IastClassLoader;
+import io.dongtai.iast.agent.IastProperties;
+import io.dongtai.iast.agent.LogCollector;
 import io.dongtai.iast.agent.fallback.FallbackManager;
 import io.dongtai.iast.agent.report.AgentRegisterReport;
-import io.dongtai.iast.agent.util.*;
+import io.dongtai.iast.agent.util.FileUtils;
+import io.dongtai.iast.agent.util.HttpClientUtils;
+import io.dongtai.iast.agent.util.ThreadUtils;
 import io.dongtai.iast.common.state.AgentState;
 import io.dongtai.log.DongTaiLog;
 import io.dongtai.log.ErrorCode;
@@ -13,6 +17,8 @@
 import java.lang.instrument.Instrumentation;
 import java.lang.management.ManagementFactory;
 import java.util.jar.JarFile;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /**
  * 引擎管理类，负责engine模块的完整生命周期，包括：下载、安装、启动、停止、重启、卸载
@@ -248,11 +254,32 @@ public synchronized boolean uninstall() {
 
     public static String getPID() {
         if (PID == null) {
-            PID = ManagementFactory.getRuntimeMXBean().getName().split("@")[0];
+            String runtimeName = ManagementFactory.getRuntimeMXBean().getName();
+            PID = extractPID(runtimeName);
         }
         return PID;
     }
 
+    /**
+     * 通过正则提取runtimeName的PID
+     * 从开头开始匹配，遇到非数字字符串停止匹配
+     * @param runtimeName 运行名称通常为PID@HostName
+     * @return PID
+     */
+    public static String extractPID(String runtimeName){
+        Pattern pattern = Pattern.compile("^\\d+");
+        Matcher matcher = pattern.matcher(runtimeName);
+
+        //防止极端情况未获取到PID ，设置默认值为0，防止服务端出现问题
+        String extractedNumber = "0";
+        if (matcher.find()) {
+            extractedNumber  = matcher.group(); // 提取匹配到的数字
+        }else {
+            DongTaiLog.warn("Get PID parsing exception, PID raw data is {}",runtimeName);
+        }
+        return extractedNumber;
+    }
+
     public AgentState getAgentState() {
         return this.agentState;
     }

dongtai-agent/src/test/java/io/dongtai/iast/agent/manager/EngineManagerTest.java

@@ -0,0 +1,22 @@
+package io.dongtai.iast.agent.manager;
+
+import io.dongtai.iast.agent.IastProperties;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.management.ManagementFactory;
+
+public class EngineManagerTest {
+
+    @Test
+    public void extractPIDTest(){
+        //初始化临时文件
+        IastProperties.initTmpDir();
+        //获取PID
+        String runtimeName = ManagementFactory.getRuntimeMXBean().getName();
+
+        //Java获取的runtimeName格式为 PID@虚拟机唯一标识 提取PID
+        String pid = EngineManager.extractPID(runtimeName);
+        Assert.assertTrue(pid.matches("\\d+"));
+    }
+}

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

@@ -35,21 +35,23 @@ public static void solveSyncInvoke(MethodEvent event, AtomicInteger invokeIdSequ
 
             Method addHeaderMethod = template.getClass().getDeclaredMethod("header", String.class, String[].class);
             addHeaderMethod.setAccessible(true);
-            String traceId = ContextManager.nextTraceId();
             // clear old traceId header
             /*
             防止高并发下的treeMap修改问题，暂时可解决
              */
             synchronized (template){
+                //将生成traceId下放到锁内
+                String traceId = ContextManager.nextTraceId();
+
                 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())});
+                event.traceId = traceId;
             }
             // add to method pool
             event.source = false;
-            event.traceId = traceId;
             event.setCallStacks(StackUtils.createCallStack(4));
             int invokeId = invokeIdSequencer.getAndIncrement();
             event.setInvokeId(invokeId);