Skip to content

Commit

Permalink
Close ISSUE HXSecurity/DongTai#147 <DongTai-agent-java>
Browse files Browse the repository at this point in the history
  • Loading branch information
owefsad committed Aug 24, 2021
1 parent 15940d0 commit 15bf796
Show file tree
Hide file tree
Showing 9 changed files with 96 additions and 39 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ public void run() {
EngineManager.turnOffLingzhi();
}
try {
Thread.sleep(properties.getReportInterval());
send();
} catch (IOException e) {
logger.error("report error reason: ", e);
Expand Down
26 changes: 26 additions & 0 deletions iast-core/src/main/java/com/secnium/iast/core/EngineManager.java
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package com.secnium.iast.core;

import com.secnium.iast.core.handler.models.IastReplayModel;
import com.secnium.iast.core.handler.models.MethodEvent;
import com.secnium.iast.core.middlewarerecognition.IastServer;
import com.secnium.iast.core.middlewarerecognition.ServerDetect;
Expand Down Expand Up @@ -37,6 +38,7 @@ public class EngineManager {

private static final ArrayBlockingQueue<String> REPORTS = new ArrayBlockingQueue<String>(2048);
private static final ArrayBlockingQueue<String> METHOD_REPORT = new ArrayBlockingQueue<String>(2048);
private final static ArrayBlockingQueue<IastReplayModel> REPLAY_QUEUE = new ArrayBlockingQueue<IastReplayModel>(256);

private static boolean logined = false;
private static int reqCounts = 0;
Expand Down Expand Up @@ -152,6 +154,26 @@ public static boolean hasNewReport() {
return !REPORTS.isEmpty();
}

public static int getReportQueueSize() {
return REPORTS.size();
}

public static boolean hasReplayData() {
return !REPLAY_QUEUE.isEmpty();
}

public static IastReplayModel getReplayModel() {
return REPLAY_QUEUE.poll();
}

public static void sendReplayModel(IastReplayModel replayModel) {
REPLAY_QUEUE.offer(replayModel);
}

public static int getReplayQueueSize() {
return REPLAY_QUEUE.size();
}

public static void sendMethodReport(String report) {
METHOD_REPORT.offer(report);
}
Expand All @@ -164,6 +186,10 @@ public static boolean hasMethodReport() {
return !METHOD_REPORT.isEmpty();
}

public static int getMethodReportQueueSize() {
return METHOD_REPORT.size();
}

public static boolean getIsLoginLogic() {
return LOGIN_LOGIC_WEIGHT.get() != null && LOGIN_LOGIC_WEIGHT.get().equals(2);
}
Expand Down
14 changes: 11 additions & 3 deletions iast-core/src/main/java/com/secnium/iast/core/PropertyUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ public class PropertyUtils {
private String iastDumpPath;
private Long heartBeatInterval = -1L;
private Long reportInterval = -1L;
private Long replayInterval = -1L;
private String serverUrl;
private String namespace;
private String engineName;
Expand Down Expand Up @@ -165,10 +166,10 @@ public boolean isEnableDumpClass() {
}

public long getReplayInterval() {
if (heartBeatInterval == -1L) {
heartBeatInterval = Long.valueOf(System.getProperty("iast.service.replay.interval", cfg.getProperty("iast.service.replay.interval", "300000")));
if (replayInterval == -1L) {
replayInterval = Long.valueOf(System.getProperty("iast.service.replay.interval", cfg.getProperty("iast.service.replay.interval", "5000")));
}
return heartBeatInterval;
return replayInterval;
}

public long getReportInterval() {
Expand All @@ -178,6 +179,13 @@ public long getReportInterval() {
return reportInterval;
}

public long getHeartBeatInterval() {
if (heartBeatInterval == -1L) {
heartBeatInterval = Long.valueOf(System.getProperty("iast.service.heartbeat.interval", cfg.getProperty("iast.service.heartbeat.interval", "5")));
}
return heartBeatInterval;
}

public String getBaseUrl() {
if (null == serverUrl) {
serverUrl = System.getProperty("iast.server.url", cfg.getProperty("iast.server.url"));
Expand Down
18 changes: 13 additions & 5 deletions iast-core/src/main/java/com/secnium/iast/core/ServiceFactory.java
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.secnium.iast.core.replay.HttpRequestReplay;
import com.secnium.iast.core.report.HeartBeatSender;
import com.secnium.iast.core.report.MethodReportSender;
import com.secnium.iast.core.report.ReportSender;

Expand All @@ -16,11 +17,14 @@ public class ServiceFactory {
private static ServiceFactory INSTANCE;
private final long replayInterval;
private final long reportInterval;
private final ScheduledExecutorService executorService;
private final long heartBeatInterval;
private final ScheduledExecutorService heartBeatService;
private final ScheduledExecutorService reportService;

ReportSender report = null;
HttpRequestReplay requestReplay = null;
MethodReportSender methodReportSender = null;
HeartBeatSender heartBeatSender = null;

public static ServiceFactory getInstance() {
if (null == INSTANCE) {
Expand All @@ -37,19 +41,23 @@ public ServiceFactory() {
PropertyUtils propertiesUtils = PropertyUtils.getInstance();
this.replayInterval = propertiesUtils.getReplayInterval();
this.reportInterval = propertiesUtils.getReportInterval();
this.executorService = Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder().setNameFormat("dongtai-engine-report").build());
this.heartBeatInterval = propertiesUtils.getHeartBeatInterval();
this.heartBeatService = Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder().setNameFormat("dongtai-heartbeat").build());
this.reportService = Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder().setNameFormat("dongtai-report").build());
}

public void init() {
heartBeatSender = new HeartBeatSender();
methodReportSender = new MethodReportSender();
report = new ReportSender();
requestReplay = new HttpRequestReplay();
}

public void start() {
executorService.scheduleWithFixedDelay(methodReportSender, 0, reportInterval, TimeUnit.MILLISECONDS);
executorService.scheduleWithFixedDelay(report, 0, reportInterval, TimeUnit.MILLISECONDS);
executorService.scheduleWithFixedDelay(requestReplay, 0, replayInterval, TimeUnit.MILLISECONDS);
heartBeatService.scheduleWithFixedDelay(heartBeatSender, 0, heartBeatInterval, TimeUnit.SECONDS);
reportService.scheduleWithFixedDelay(methodReportSender, 0, reportInterval, TimeUnit.MILLISECONDS);
reportService.scheduleWithFixedDelay(report, 0, reportInterval, TimeUnit.MILLISECONDS);
reportService.scheduleWithFixedDelay(requestReplay, 0, replayInterval, TimeUnit.MILLISECONDS);
}

public void stop() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,4 +86,7 @@ public class ReportConstant {
public static final String OVER_POWER_AUTH_COOKIE = "cookie";
public static final String OVER_POWER_TRACE_ID = "x-trace-id";

public static final String REPORT_QUEUE = "report_queue";
public static final String METHOD_QUEUE = "method_queue";
public static final String REPLAY_QUEUE = "replay_queue";
}
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package com.secnium.iast.core.replay;

import com.secnium.iast.core.AbstractThread;
import com.secnium.iast.core.EngineManager;
import com.secnium.iast.core.handler.models.IastReplayModel;
import com.secnium.iast.core.report.ErrorLogReport;
import com.secnium.iast.core.util.HttpClientHostnameVerifier;
Expand All @@ -22,7 +23,6 @@
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;

/**
* 根据字符串格式的header头及cookie信息生成重放请求
Expand All @@ -39,11 +39,7 @@
public class HttpRequestReplay extends AbstractThread {
private static final String PROTOCOL_HTTPS = "https";
public final static HostnameVerifier DO_NOT_VERIFY = new HttpClientHostnameVerifier();
private final static ArrayBlockingQueue<IastReplayModel> WAITING_REPLAY_REQUESTS = new ArrayBlockingQueue<IastReplayModel>(256);

public static void sendReplayRequest(IastReplayModel replayModel) {
WAITING_REPLAY_REQUESTS.offer(replayModel);
}

/**
* @param replayRequestRaw
Expand All @@ -66,7 +62,7 @@ public static void sendReplayRequest(StringBuilder replayRequestRaw) {
replayRequest.get("relation_id"),
replayRequest.get("replay_type"));
if (replayModel.isValid()) {
sendReplayRequest(replayModel);
EngineManager.sendReplayModel(replayModel);
}
}
}
Expand Down Expand Up @@ -171,8 +167,8 @@ private static void sendRequest(String method, String fullUrl, String data, Hash

@Override
protected void send() {
while (!WAITING_REPLAY_REQUESTS.isEmpty()) {
IastReplayModel model = WAITING_REPLAY_REQUESTS.poll();
while (EngineManager.hasReplayData()) {
IastReplayModel model = EngineManager.getReplayModel();
doReplay(model);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,15 @@

import com.secnium.iast.core.EngineManager;
import com.secnium.iast.core.handler.vulscan.ReportConstant;
import com.secnium.iast.core.replay.HttpRequestReplay;
import com.secnium.iast.core.util.ByteUtils;
import com.secnium.iast.core.util.Constants;
import com.secnium.iast.core.util.HttpClientUtils;
import com.secnium.iast.core.util.LogUtils;
import org.json.JSONObject;
import org.slf4j.Logger;

import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryUsage;
import java.lang.management.OperatingSystemMXBean;
Expand All @@ -14,12 +20,8 @@
*
* @author [email protected]
*/
public class HeartBeatReport {

public static void createReport() {
String msg = generateHeartBeatMsg();
EngineManager.sendNewReport(msg);
}
public class HeartBeatSender extends Thread {
private final Logger logger = LogUtils.getLogger(getClass());

private static String generateHeartBeatMsg() {
JSONObject report = new JSONObject();
Expand All @@ -32,6 +34,9 @@ private static String generateHeartBeatMsg() {
detail.put(ReportConstant.HEART_BEAT_CPU, readCpuInfo());
detail.put(ReportConstant.HEART_BEAT_DISK, getDiskInfo());
detail.put(ReportConstant.HEART_BEAT_REQ_COUNT, EngineManager.getRequestCount());
detail.put(ReportConstant.REPORT_QUEUE, EngineManager.getReportQueueSize());
detail.put(ReportConstant.METHOD_QUEUE, EngineManager.getMethodReportQueueSize());
detail.put(ReportConstant.REPLAY_QUEUE, EngineManager.getReplayQueueSize());

return report.toString();
}
Expand Down Expand Up @@ -68,4 +73,25 @@ public static String getMemInfo() {
public static String getDiskInfo() {
return "{}";
}

@Override
public void run() {
boolean isRunning = EngineManager.isLingzhiRunning();
if (isRunning) {
EngineManager.turnOffLingzhi();
}
try {
String report = generateHeartBeatMsg();
StringBuilder response = HttpClientUtils.sendPost(Constants.API_REPORT_UPLOAD, report);
HttpRequestReplay.sendReplayRequest(response);
} catch (IOException e) {
logger.error("report error reason: ", e);
} catch (Exception e) {
logger.error("report error, reason: ", e);
}

if (isRunning) {
EngineManager.turnOnLingzhi();
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,40 +2,30 @@

import com.secnium.iast.core.AbstractThread;
import com.secnium.iast.core.EngineManager;
import com.secnium.iast.core.replay.HttpRequestReplay;
import com.secnium.iast.core.util.Constants;
import com.secnium.iast.core.util.HttpClientUtils;
import com.secnium.iast.core.util.LogUtils;
import org.slf4j.Logger;

import java.util.regex.Pattern;

/**
* 发送报告的功能实现
*
* @author [email protected]
*/
public class ReportSender extends AbstractThread {
private final Logger logger = LogUtils.getLogger(ReportSender.class);
private final Pattern PATTERN = Pattern.compile("\"type\":1}");
// private final Pattern PATTERN = Pattern.compile("\"type\":1}");

@Override
protected void send() throws Exception {
HeartBeatReport.createReport();
String report;
StringBuilder response;
while (EngineManager.hasNewReport()) {
report = EngineManager.getNewReport();
try {
if (report != null && !report.isEmpty()) {
response = HttpClientUtils.sendPost(Constants.API_REPORT_UPLOAD, report);
if (PATTERN.matcher(report).find()) {
HttpRequestReplay.sendReplayRequest(response);
}
}
} catch (Exception e) {
logger.info(report);
throw e;
if (logger.isDebugEnabled()) {
logger.debug(report);
}
if (report != null && !report.isEmpty()) {
HttpClientUtils.sendPost(Constants.API_REPORT_UPLOAD, report);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ public static Object cloneResponse(Object response) {

public ResponseWrapper(HttpServletResponse response) {
super(response);
response.addHeader("DongTai", "v1.0.3");
}

@Override
Expand Down

0 comments on commit 15bf796

Please sign in to comment.