Skip to content

Commit

Permalink
Merge pull request #513 from Nizernizer/fature/allow-data-report
Browse files Browse the repository at this point in the history 
feature: allow data report
  • Loading branch information
@lostsnow
lostsnow authored Apr 26, 2023
2 parents 901d77e + 8f743b2 commit d61c0ed
Show file tree
Hide file tree
Showing 7 changed files with 53 additions and 36 deletions.
9 changes: 1 addition & 8 deletions dongtai-agent/src/main/java/io/dongtai/iast/agent/AgentLauncher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,13 +164,6 @@ private static void install(final Instrumentation inst) {
if (send) {
LogCollector.extractFluent();
DongTaiLog.info("Agent registered successfully.");
Boolean agentStat = AgentRegisterReport.agentStat();
if (!agentStat) {
AgentStateMonitor.isCoreRegisterStart = false;
DongTaiLog.info("Detection engine not started, agent waiting to be audited.");
} else {
AgentStateMonitor.isCoreRegisterStart = true;
}
shutdownHook = new ShutdownThread();
Runtime.getRuntime().addShutdownHook(shutdownHook);
loadEngine(inst);
Expand All @@ -187,7 +180,7 @@ private static void install(final Instrumentation inst) {
private static void loadEngine(final Instrumentation inst) {
EngineManager engineManager = EngineManager.getInstance(inst, LAUNCH_MODE, EngineManager.getPID(), AGENT_STATE);
MonitorDaemonThread daemonThread = MonitorDaemonThread.getInstance(engineManager);
if (MonitorDaemonThread.delayTime <= 0 && AgentStateMonitor.isCoreRegisterStart) {
if (MonitorDaemonThread.delayTime <= 0) {
daemonThread.startEngine();
}

Expand Down
3 changes: 2 additions & 1 deletion dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import io.dongtai.iast.agent.*;
import io.dongtai.iast.agent.fallback.FallbackManager;
import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
import io.dongtai.iast.agent.report.AgentRegisterReport;
import io.dongtai.iast.agent.util.*;
import io.dongtai.iast.common.state.AgentState;
Expand Down Expand Up @@ -36,7 +37,7 @@ public class EngineManager {
private final IastProperties properties;
private final String launchMode;
private Class<?> classOfEngine;
private FallbackManager fallbackManager;
private final FallbackManager fallbackManager;
private final AgentState agentState;

/**
Expand Down
5 changes: 1 addition & 4 deletions dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/MonitorDaemonThread.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,8 @@ public void run() {
if (MonitorDaemonThread.delayTime > 0) {
try {
Thread.sleep(delayTime);
} catch (InterruptedException ignore) {
}
if (AgentStateMonitor.isCoreRegisterStart) {
startEngine();
} catch (InterruptedException ignore) {
}
}
// 引擎启动成功后，创建子线程执行monitor任务
Expand All @@ -81,7 +79,6 @@ public void startEngine() {
// jdk8以上
status = engineManager.extractPackage();
status = status && engineManager.install();
status = status && engineManager.start();
}
if (!status) {
DongTaiLog.info("DongTai IAST started failure");
Expand Down
55 changes: 39 additions & 16 deletions dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/AgentStateMonitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import io.dongtai.iast.agent.util.ThreadUtils;
import io.dongtai.iast.common.constants.AgentConstant;
import io.dongtai.iast.common.constants.ApiPath;
import io.dongtai.iast.common.state.AgentState;
import io.dongtai.iast.common.state.State;
import io.dongtai.iast.common.state.StateCause;
import io.dongtai.log.DongTaiLog;
Expand All @@ -24,7 +25,6 @@
*/
public class AgentStateMonitor implements IMonitor {
private final EngineManager engineManager;
public static Boolean isCoreRegisterStart = false;
private static final String NAME = "AgentStateMonitor";

public AgentStateMonitor(EngineManager engineManager) {
Expand All @@ -38,50 +38,73 @@ public String getName() {

@Override
public void check() {
AgentState agentState = this.engineManager.getAgentState();
try {
if (this.engineManager.getAgentState().getState() == null) {
if (agentState.getState() == null) {
return;
}

if (this.engineManager.getAgentState().isUninstalledByCli()) {
if (agentState.isUninstalledByCli()) {
HttpClientUtils.sendPost(ApiPath.ACTUAL_ACTION,
HeartBeatReport.generateAgentActualActionMsg(this.engineManager.getAgentState()));
HeartBeatReport.generateAgentActualActionMsg(agentState));
return;
}

if (!this.engineManager.getAgentState().isFallback() && !this.engineManager.getAgentState().isException()) {
String expectState = checkExpectState();
if (State.RUNNING.equals(expectState) && this.engineManager.getAgentState().isPaused()) {
Map<String, String> stringStringMap = checkExpectState();
// 默认值
String expectState = "other";
boolean allowReport = true;

if (stringStringMap != null) {
expectState = stringStringMap.get("exceptRunningStatus");
if (null != stringStringMap.get("isAllowDateReport")) {
allowReport = !"0".equals(stringStringMap.get("isAllowDateReport"));
}
}

if (allowReport && !agentState.isAllowReport()) {
DongTaiLog.info("engine is allowed to report data");
agentState.setAllowReport(allowReport);
} else if (!allowReport && agentState.isAllowReport()) {
DongTaiLog.info("engine is not allowed to report data");
agentState.setAllowReport(allowReport);
}

if (!agentState.isFallback() && !agentState.isException() && agentState.isAllowReport() && agentState.isAllowReport()) {
if (State.RUNNING.equals(expectState) && agentState.isPaused()) {
DongTaiLog.info("engine start by server expect state");
engineManager.start();
engineManager.getAgentState().setState(State.RUNNING).setCause(StateCause.RUNNING_BY_SERVER);
} else if (State.PAUSED.equals(expectState) && this.engineManager.getAgentState().isRunning()) {
agentState.setState(State.RUNNING).setCause(StateCause.RUNNING_BY_SERVER);
} else if (State.PAUSED.equals(expectState) && agentState.isRunning()) {
DongTaiLog.info("engine stop by server expect state");
engineManager.stop();
engineManager.getAgentState().setState(State.PAUSED).setCause(StateCause.PAUSE_BY_SERVER);
agentState.setState(State.PAUSED).setCause(StateCause.PAUSE_BY_SERVER);
}
}
HttpClientUtils.sendPost(ApiPath.ACTUAL_ACTION,
HeartBeatReport.generateAgentActualActionMsg(this.engineManager.getAgentState()));
HeartBeatReport.generateAgentActualActionMsg(agentState));
} catch (Throwable t) {
DongTaiLog.warn(ErrorCode.AGENT_MONITOR_THREAD_CHECK_FAILED, getName(), t);
}
}

private String checkExpectState() {
private Map<String, String> checkExpectState() {
try {
Map<String, String> parameters = new HashMap<String, String>();
Map<String, String> parameters = new HashMap<>();
parameters.put("agentId", String.valueOf(AgentRegisterReport.getAgentId()));
String respRaw = HttpClientUtils.sendGet(ApiPath.EXCEPT_ACTION, parameters).toString();
if (!respRaw.isEmpty()) {
JSONObject resp = JSON.parseObject(respRaw);
JSONObject data = (JSONObject) resp.get("data");
return data.get("exceptRunningStatus").toString();
Map<String, String> objectObjectHashMap = new HashMap<>(2);
String s = data.toJSONString();
objectObjectHashMap = JSON.parseObject(s, Map.class);
return objectObjectHashMap;
}
} catch (Throwable e) {
return "other";
return null;
}
return "other";
return null;
}

@Override
Expand Down
6 changes: 0 additions & 6 deletions dongtai-agent/src/main/java/io/dongtai/iast/agent/report/AgentRegisterReport.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ public class AgentRegisterReport {
public static AgentRegisterReport INSTANCE;
private String projectName = null;
private static Integer agentId = -1;
private static Integer coreRegisterStart = 1;
final IServer server = ServerDetect.getWebserver();
private static String AGENT_NAME = null;
private static String HOST_NAME = null;
Expand Down Expand Up @@ -265,7 +264,6 @@ private void setAgentData(StringBuilder responseRaw) {
if (status == 201) {
JSONObject data = (JSONObject) responseObj.get("data");
agentId = (Integer) data.get("id");
coreRegisterStart = (Integer) data.get("coreAutoStart");
} else {
DongTaiLog.error(ErrorCode.AGENT_REGISTER_RESPONSE_CODE_INVALID, responseRaw);
}
Expand All @@ -275,10 +273,6 @@ private void setAgentData(StringBuilder responseRaw) {
}
}

public static Boolean agentStat() {
return coreRegisterStart == 1;
}

private static String generateUUID() {
String uuidPath = IastProperties.getInstance().getUUIDPath();
if (uuidPath == null || uuidPath.isEmpty()) {
Expand Down
9 changes: 9 additions & 0 deletions dongtai-common/src/main/java/io/dongtai/iast/common/state/AgentState.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ public class AgentState {
private State pendingState;
private StateCause cause;
private boolean fallback;
private boolean allowReport = true;
private static AgentState INSTANCE;

public static AgentState getInstance() {
Expand Down Expand Up @@ -95,4 +96,12 @@ public void fallbackRecover() {
public boolean isFallback() {
return this.fallback;
}

public boolean isAllowReport() {
return this.allowReport;
}

public void setAllowReport(boolean allowReport) {
this.allowReport = allowReport;
}
}
2 changes: 1 addition & 1 deletion dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ public static int getRequestCount() {
* @return true - 引擎已启动；false - 引擎未启动
*/
public static boolean isEngineRunning() {
return AGENT_STATE.isRunning() && AGENT_STATE.getPendingState() == null;
return AGENT_STATE.isRunning() && AGENT_STATE.getPendingState() == null && AGENT_STATE.isAllowReport();
}

public boolean isEnableDumpClass() {
Expand Down

0 comments on commit d61c0ed

Please sign in to comment.