Skip to content

Commit 6c15feb

Browse files
author
‘niuerzhuang’
committed
feature: allow data report
1 parent 5f8ecc8 commit 6c15feb

File tree

7 files changed

+73
-33
lines changed

7 files changed

+73
-33
lines changed

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/AgentLauncher.java‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,8 @@ public static void agentmain(String args, Instrumentation inst) {
152152
public static synchronized void uninstall() {
153153
EngineManager engineManager = EngineManager.getInstance();
154154
engineManager.uninstall();
155+
// 手动卸载时将线程停止
156+
MonitorDaemonThread.isExit = true;
155157
}
156158

157159
/**
@@ -164,13 +166,6 @@ private static void install(final Instrumentation inst) {
164166
if (send) {
165167
LogCollector.extractFluent();
166168
DongTaiLog.info("Agent registered successfully.");
167-
Boolean agentStat = AgentRegisterReport.agentStat();
168-
if (!agentStat) {
169-
AgentStateMonitor.isCoreRegisterStart = false;
170-
DongTaiLog.info("Detection engine not started, agent waiting to be audited.");
171-
} else {
172-
AgentStateMonitor.isCoreRegisterStart = true;
173-
}
174169
shutdownHook = new ShutdownThread();
175170
Runtime.getRuntime().addShutdownHook(shutdownHook);
176171
loadEngine(inst);
@@ -187,7 +182,7 @@ private static void install(final Instrumentation inst) {
187182
private static void loadEngine(final Instrumentation inst) {
188183
EngineManager engineManager = EngineManager.getInstance(inst, LAUNCH_MODE, EngineManager.getPID(), AGENT_STATE);
189184
MonitorDaemonThread daemonThread = MonitorDaemonThread.getInstance(engineManager);
190-
if (MonitorDaemonThread.delayTime <= 0 && AgentStateMonitor.isCoreRegisterStart) {
185+
if (MonitorDaemonThread.delayTime <= 0) {
191186
daemonThread.startEngine();
192187
}
193188

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io.dongtai.iast.agent.*;
44
import io.dongtai.iast.agent.fallback.FallbackManager;
5+
import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
56
import io.dongtai.iast.agent.report.AgentRegisterReport;
67
import io.dongtai.iast.agent.util.*;
78
import io.dongtai.iast.common.state.AgentState;
@@ -36,7 +37,7 @@ public class EngineManager {
3637
private final IastProperties properties;
3738
private final String launchMode;
3839
private Class<?> classOfEngine;
39-
private FallbackManager fallbackManager;
40+
private final FallbackManager fallbackManager;
4041
private final AgentState agentState;
4142

4243
/**
@@ -175,6 +176,7 @@ public boolean install() {
175176
String.class)
176177
.invoke(null, launchMode, this.properties.getPropertiesFilePath(),
177178
AgentRegisterReport.getAgentId(), inst, agentPath);
179+
MonitorDaemonThread.isExit = false;
178180
return true;
179181
} catch (Throwable e) {
180182
DongTaiLog.error(ErrorCode.AGENT_REFLECTION_INSTALL_FAILED, e);

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/MonitorDaemonThread.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,8 @@ public void run() {
5555
if (MonitorDaemonThread.delayTime > 0) {
5656
try {
5757
Thread.sleep(delayTime);
58-
} catch (InterruptedException ignore) {
59-
}
60-
if (AgentStateMonitor.isCoreRegisterStart) {
6158
startEngine();
59+
} catch (InterruptedException ignore) {
6260
}
6361
}
6462
// 引擎启动成功后，创建子线程执行monitor任务
@@ -81,7 +79,6 @@ public void startEngine() {
8179
// jdk8以上
8280
status = engineManager.extractPackage();
8381
status = status && engineManager.install();
84-
status = status && engineManager.start();
8582
}
8683
if (!status) {
8784
DongTaiLog.info("DongTai IAST started failure");

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/AgentStateMonitor.java‎

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import io.dongtai.iast.agent.util.ThreadUtils;
1212
import io.dongtai.iast.common.constants.AgentConstant;
1313
import io.dongtai.iast.common.constants.ApiPath;
14+
import io.dongtai.iast.common.state.AgentState;
1415
import io.dongtai.iast.common.state.State;
1516
import io.dongtai.iast.common.state.StateCause;
1617
import io.dongtai.log.DongTaiLog;
@@ -24,7 +25,6 @@
2425
*/
2526
public class AgentStateMonitor implements IMonitor {
2627
private final EngineManager engineManager;
27-
public static Boolean isCoreRegisterStart = false;
2828
private static final String NAME = "AgentStateMonitor";
2929

3030
public AgentStateMonitor(EngineManager engineManager) {
@@ -38,50 +38,81 @@ public String getName() {
3838

3939
@Override
4040
public void check() {
41+
AgentState agentState = this.engineManager.getAgentState();
4142
try {
42-
if (this.engineManager.getAgentState().getState() == null) {
43+
if (agentState.getState() == null) {
4344
return;
4445
}
4546

46-
if (this.engineManager.getAgentState().isUninstalledByCli()) {
47+
if (agentState.isUninstalledByCli()) {
4748
HttpClientUtils.sendPost(ApiPath.ACTUAL_ACTION,
48-
HeartBeatReport.generateAgentActualActionMsg(this.engineManager.getAgentState()));
49+
HeartBeatReport.generateAgentActualActionMsg(agentState));
4950
return;
5051
}
5152

52-
if (!this.engineManager.getAgentState().isFallback() && !this.engineManager.getAgentState().isException()) {
53-
String expectState = checkExpectState();
54-
if (State.RUNNING.equals(expectState) && this.engineManager.getAgentState().isPaused()) {
53+
Map<String, String> stringStringMap = checkExpectState();
54+
// 默认值
55+
String expectState = "other";
56+
String allowReport = "1";
57+
58+
if (stringStringMap != null) {
59+
expectState = stringStringMap.get("exceptRunningStatus");
60+
allowReport = stringStringMap.get("isAllowDateReport");
61+
}
62+
63+
if (!agentState.isAllowReport(allowReport)) {
64+
if (null == agentState.getAllowReport()) {
65+
DongTaiLog.info("engine is not allowed to report data");
66+
agentState.setAllowReport(allowReport);
67+
return;
68+
}
69+
if (!allowReport.equals(agentState.getAllowReport())) {
70+
DongTaiLog.info("engine is not allowed to report data");
71+
agentState.setAllowReport(allowReport);
72+
engineManager.stop();
73+
}
74+
} else if (agentState.isAllowReport(allowReport) && !allowReport.equals(agentState.getAllowReport())) {
75+
DongTaiLog.info("engine is allowed to report data");
76+
agentState.setAllowReport(allowReport);
77+
engineManager.start();
78+
return;
79+
}
80+
81+
if (!agentState.isFallback() && !agentState.isException()) {
82+
if (State.RUNNING.equals(expectState) && agentState.isPaused()) {
5583
DongTaiLog.info("engine start by server expect state");
5684
engineManager.start();
5785
engineManager.getAgentState().setState(State.RUNNING).setCause(StateCause.RUNNING_BY_SERVER);
58-
} else if (State.PAUSED.equals(expectState) && this.engineManager.getAgentState().isRunning()) {
86+
} else if (State.PAUSED.equals(expectState) && agentState.isRunning()) {
5987
DongTaiLog.info("engine stop by server expect state");
6088
engineManager.stop();
6189
engineManager.getAgentState().setState(State.PAUSED).setCause(StateCause.PAUSE_BY_SERVER);
6290
}
6391
}
6492
HttpClientUtils.sendPost(ApiPath.ACTUAL_ACTION,
65-
HeartBeatReport.generateAgentActualActionMsg(this.engineManager.getAgentState()));
93+
HeartBeatReport.generateAgentActualActionMsg(agentState));
6694
} catch (Throwable t) {
6795
DongTaiLog.warn(ErrorCode.AGENT_MONITOR_THREAD_CHECK_FAILED, getName(), t);
6896
}
6997
}
7098

71-
private String checkExpectState() {
99+
private Map<String, String> checkExpectState() {
72100
try {
73101
Map<String, String> parameters = new HashMap<String, String>();
74102
parameters.put("agentId", String.valueOf(AgentRegisterReport.getAgentId()));
75103
String respRaw = HttpClientUtils.sendGet(ApiPath.EXCEPT_ACTION, parameters).toString();
76104
if (!respRaw.isEmpty()) {
77105
JSONObject resp = JSON.parseObject(respRaw);
78106
JSONObject data = (JSONObject) resp.get("data");
79-
return data.get("exceptRunningStatus").toString();
107+
Map<String, String> objectObjectHashMap = new HashMap<>(2);
108+
String s = data.toJSONString();
109+
objectObjectHashMap = JSON.parseObject(s, Map.class);
110+
return objectObjectHashMap;
80111
}
81112
} catch (Throwable e) {
82-
return "other";
113+
return null;
83114
}
84-
return "other";
115+
return null;
85116
}
86117

87118
@Override

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/report/AgentRegisterReport.java‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ public class AgentRegisterReport {
2929
public static AgentRegisterReport INSTANCE;
3030
private String projectName = null;
3131
private static Integer agentId = -1;
32-
private static Integer coreRegisterStart = 1;
3332
final IServer server = ServerDetect.getWebserver();
3433
private static String AGENT_NAME = null;
3534
private static String HOST_NAME = null;
@@ -265,7 +264,6 @@ private void setAgentData(StringBuilder responseRaw) {
265264
if (status == 201) {
266265
JSONObject data = (JSONObject) responseObj.get("data");
267266
agentId = (Integer) data.get("id");
268-
coreRegisterStart = (Integer) data.get("coreAutoStart");
269267
} else {
270268
DongTaiLog.error(ErrorCode.AGENT_REGISTER_RESPONSE_CODE_INVALID, responseRaw);
271269
}
@@ -275,10 +273,6 @@ private void setAgentData(StringBuilder responseRaw) {
275273
}
276274
}
277275

278-
public static Boolean agentStat() {
279-
return coreRegisterStart == 1;
280-
}
281-
282276
private static String generateUUID() {
283277
String uuidPath = IastProperties.getInstance().getUUIDPath();
284278
if (uuidPath == null || uuidPath.isEmpty()) {

‎dongtai-common/src/main/java/io/dongtai/iast/common/state/AgentState.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ public class AgentState {
55
private State pendingState;
66
private StateCause cause;
77
private boolean fallback;
8+
private Integer allowReport;
89
private static AgentState INSTANCE;
910

1011
public static AgentState getInstance() {
@@ -95,4 +96,24 @@ public void fallbackRecover() {
9596
public boolean isFallback() {
9697
return this.fallback;
9798
}
99+
100+
public Boolean isAllowReport(String allowReport) {
101+
if (null == allowReport) {
102+
return true;
103+
} else {
104+
return Integer.valueOf(allowReport).equals(1);
105+
}
106+
}
107+
108+
public Boolean isAllowReport() {
109+
return 1 == this.allowReport;
110+
}
111+
112+
public void setAllowReport(String allowReport) {
113+
this.allowReport = Integer.valueOf(allowReport);
114+
}
115+
116+
public String getAllowReport() {
117+
return String.valueOf(allowReport);
118+
}
98119
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ public static int getRequestCount() {
8989
* @return true - 引擎已启动；false - 引擎未启动
9090
*/
9191
public static boolean isEngineRunning() {
92-
return AGENT_STATE.isRunning() && AGENT_STATE.getPendingState() == null;
92+
return AGENT_STATE.isRunning() && AGENT_STATE.getPendingState() == null && AGENT_STATE.isAllowReport();
9393
}
9494

9595
public boolean isEnableDumpClass() {

0 commit comments

Comments
 (0)