Skip to content

Commit

Permalink
Merge pull request #334 from Nizernizer/main
Browse files Browse the repository at this point in the history 
fix: agent upgrade and resource load.
  • Loading branch information
@lostsnow
lostsnow authored Jul 13, 2022
2 parents 2de6a23 + e6eb23e commit 8417960
Show file tree
Hide file tree
Showing 7 changed files with 90 additions and 64 deletions.
22 changes: 11 additions & 11 deletions dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,9 @@

import java.io.*;
import java.util.Arrays;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

import io.dongtai.iast.agent.util.FileUtils;
import io.dongtai.log.DongTaiLog;
import org.apache.commons.cli.*;

/**
Expand Down Expand Up @@ -103,14 +102,14 @@ private static void doAttach(String pid, String agentArgs) {
Process process = Runtime.getRuntime().exec(execution);
process.waitFor();
if (process.exitValue() == 0) {
DongTaiLog.info("attach to process {} success, command: {}", pid, Arrays.toString(execution));
System.out.println("[io.dongtai.iast.agent] [INFO] attach to process "+pid+" success, command: "+Arrays.toString(execution));
} else {
DongTaiLog.error("attach failure, please try again with command: {}", Arrays.toString(execution));
System.out.println("[io.dongtai.iast.agent] [ERROR] attach failure, please try again with command: "+Arrays.toString(execution));
}
} catch (IOException e) {
DongTaiLog.error("io.dongtai.iast.agent.Agent.doAttach(java.lang.String,java.lang.String)",e);
e.printStackTrace();
} catch (InterruptedException e) {
DongTaiLog.error("io.dongtai.iast.agent.Agent.doAttach(java.lang.String,java.lang.String)",e);
e.printStackTrace();
}
}

Expand Down Expand Up @@ -141,9 +140,9 @@ private static void extractJattach() throws IOException {
FileUtils.getResourceToFile("bin/jattach-linux", JATTACH_FILE);
}
if ((new File(JATTACH_FILE)).setExecutable(true)) {
DongTaiLog.info("jattach extract success. wait for attach");
System.out.println("[io.dongtai.iast.agent] [INFO] jattach extract success. wait for attach");
} else {
DongTaiLog.info("jattach extract failure. please set execute permission, file: {}", JATTACH_FILE);
System.out.println("[io.dongtai.iast.agent] [INFO] jattach extract failure. please set execute permission, file: "+JATTACH_FILE);
}
}

Expand All @@ -152,7 +151,8 @@ private static void extractJattach() throws IOException {
*
* @param args
*/
public static void main(String[] args) {
public static void main(String[] args) throws InterruptedException {
TimeUnit.SECONDS.sleep(10);
String[] agentArgs = new String[0];
try {
agentArgs = parseAgentArgs(args);
Expand All @@ -162,9 +162,9 @@ public static void main(String[] args) {
doAttach(agentArgs[0], agentArgs[1]);
}
} catch (ParseException e) {
DongTaiLog.error(e);
e.printStackTrace();
} catch (IOException e) {
DongTaiLog.error(e);
e.printStackTrace();
}
}

Expand Down
83 changes: 44 additions & 39 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 @@ -79,13 +79,55 @@ public static void premain(String args, Instrumentation inst) {
* @param inst inst
*/
public static void agentmain(String args, Instrumentation inst) {
DongTaiLog.info("Protect By DongTai IAST: " + System.getProperty("protect.by.dongtai", "false"));
Map<String, String> argsMap = parseArgs(args);
try {
if (argsMap.containsKey("debug")) {
System.setProperty("dongtai.debug", argsMap.get("debug"));
}
if (argsMap.containsKey("appCreate")) {
System.setProperty("dongtai.app.create", argsMap.get("appCreate"));
}
if (argsMap.containsKey("appName")) {
System.setProperty("dongtai.app.name", argsMap.get("appName"));
}
if (argsMap.containsKey("appVersion")) {
System.setProperty("dongtai.app.version", argsMap.get("appVersion"));
}
if (argsMap.containsKey("clusterName")) {
System.setProperty("dongtai.cluster.name", argsMap.get("clusterName"));
}
if (argsMap.containsKey("clusterVersion")) {
System.setProperty("dongtai.cluster.version", argsMap.get("clusterVersion"));
}
if (argsMap.containsKey("dongtaiServer")) {
System.setProperty("dongtai.server.url", argsMap.get("dongtaiServer"));
}
if (argsMap.containsKey("dongtaiToken")) {
System.setProperty("dongtai.server.token", argsMap.get("dongtaiToken"));
}
if (argsMap.containsKey("serverPackage")) {
System.setProperty("dongtai.server.package", argsMap.get("serverPackage"));
}
if (argsMap.containsKey("logLevel")) {
System.setProperty("dongtai.log.level", argsMap.get("logLevel"));
}
if (argsMap.containsKey("logPath")) {
System.setProperty("dongtai.log.path", argsMap.get("logPath"));
}
} catch (Exception e) {
DongTaiLog.error(e);
}
String tmpdir = System.getProperty("java.io.tmpdir");
String appName = System.getProperty("dongtai.app.name");
String appVersion = System.getProperty("dongtai.app.version");
System.setProperty("java.io.tmpdir.dongtai", tmpdir + File.separator + appName + "-" + appVersion + "-" + UUID.randomUUID().toString().replaceAll("-", "") + File.separator);
DongTaiLog.info("Protect By DongTai IAST: " + System.getProperty("protect.by.dongtai", "false"));
if ("uninstall".equals(argsMap.get("mode"))) {
if (System.getProperty("protect.by.dongtai", null) == null) {
DongTaiLog.info("DongTai wasn't installed.");
return;
}
EngineMonitor.setIsUninstallHeart(true);
DongTaiLog.info("Engine is about to be uninstalled");
uninstall();
// attach手动卸载后停止守护线程
Expand All @@ -98,44 +140,7 @@ public static void agentmain(String args, Instrumentation inst) {
}
MonitorDaemonThread.isExit = false;
LAUNCH_MODE = LAUNCH_MODE_ATTACH;
try {
if (argsMap.containsKey("debug")) {
System.setProperty("dongtai.debug", argsMap.get("debug"));
}
if (argsMap.containsKey("appCreate")) {
System.setProperty("dongtai.app.create", argsMap.get("appCreate"));
}
if (argsMap.containsKey("appName")) {
System.setProperty("dongtai.app.name", argsMap.get("appName"));
}
if (argsMap.containsKey("appVersion")) {
System.setProperty("dongtai.app.version", argsMap.get("appVersion"));
}
if (argsMap.containsKey("clusterName")) {
System.setProperty("dongtai.cluster.name", argsMap.get("clusterName"));
}
if (argsMap.containsKey("clusterVersion")) {
System.setProperty("dongtai.cluster.version", argsMap.get("clusterVersion"));
}
if (argsMap.containsKey("dongtaiServer")) {
System.setProperty("dongtai.server.url", argsMap.get("dongtaiServer"));
}
if (argsMap.containsKey("dongtaiToken")) {
System.setProperty("dongtai.server.token", argsMap.get("dongtaiToken"));
}
if (argsMap.containsKey("serverPackage")) {
System.setProperty("dongtai.server.package", argsMap.get("serverPackage"));
}
if (argsMap.containsKey("logLevel")) {
System.setProperty("dongtai.log.level", argsMap.get("logLevel"));
}
if (argsMap.containsKey("logPath")) {
System.setProperty("dongtai.log.path", argsMap.get("logPath"));
}
install(inst);
} catch (Exception e) {
DongTaiLog.error(e);
}
install(inst);
}
}

Expand Down
7 changes: 5 additions & 2 deletions 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 @@ -5,6 +5,7 @@
import io.dongtai.iast.agent.middlewarerecognition.ServerDetect;
import io.dongtai.iast.agent.middlewarerecognition.tomcat.AbstractTomcat;
import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
import io.dongtai.iast.agent.monitor.impl.EngineMonitor;
import io.dongtai.iast.agent.monitor.impl.PerformanceMonitor;
import io.dongtai.iast.agent.report.AgentRegisterReport;
import io.dongtai.iast.agent.util.FileUtils;
Expand Down Expand Up @@ -449,10 +450,12 @@ public synchronized boolean uninstall() {
classOfEngine = null;
IAST_CLASS_LOADER.closeIfPossible();
IAST_CLASS_LOADER = null;
uninstallObject();
setRunningStatus(1);
setCoreStop(true);
MonitorDaemonThread.isExit = true;
if (EngineMonitor.getIsUninstallHeart()){
uninstallObject();
MonitorDaemonThread.isExit = true;
}
return true;
}

Expand Down
11 changes: 10 additions & 1 deletion dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/EngineMonitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ public class EngineMonitor implements IMonitor {
private final EngineManager engineManager;
public static Boolean isCoreRegisterStart = false;
private static final String NAME = "EngineMonitor";

private static Boolean isUninstallHeart = true;

public EngineMonitor(EngineManager engineManager) {
this.engineManager = engineManager;
Expand Down Expand Up @@ -58,6 +58,7 @@ public void check() throws Exception {
break;
case CORE_UNINSTALL:
DongTaiLog.info("engine uninstall");
setIsUninstallHeart(false);
engineManager.uninstall();
break;
case CORE_PERFORMANCE_FORCE_OPEN:
Expand Down Expand Up @@ -125,6 +126,14 @@ private boolean couldInstallEngine() {
return true;
}

public static Boolean getIsUninstallHeart() {
return isUninstallHeart;
}

public static void setIsUninstallHeart(Boolean isUninstallHeart) {
EngineMonitor.isUninstallHeart = isUninstallHeart;
}

@Override
public void run() {
try {
Expand Down
2 changes: 1 addition & 1 deletion .../core/bytecode/enhance/plugin/framework/j2ee/dispatch/ServletDispatcherAdviceAdapter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ protected void before() {
mv.visitJumpInsn(EQ, elseLabel);

cloneHttpServletRequest();
cloneHttpServletResponse();
captureMethodState(-1, HookType.HTTP.getValue(), false);
cloneHttpServletResponse();
mark(elseLabel);
}

Expand Down
27 changes: 18 additions & 9 deletions ...i-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/HttpImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ public class HttpImpl {
private static Class<?> CLASS_OF_SERVLET_PROXY;
private static IastClassLoader iastClassLoader;
public static File IAST_REQUEST_JAR_PACKAGE;
private final static ThreadLocal<Map<String, Object>> REQUEST_META = new ThreadLocal<>();

static {
IAST_REQUEST_JAR_PACKAGE = new File(System.getProperty("java.io.tmpdir.dongtai") + "iast" + File.separator + "dongtai-api.jar");
Expand Down Expand Up @@ -102,10 +103,16 @@ public static Object cloneRequest(Object req, boolean isJakarta) {
* @return dongtai response object
*/
public static Object cloneResponse(Object response, boolean isJakarta) {
if (response == null) {
return null;
}
try {
if (response == null) {
return null;
}
if (ConfigMatcher.getInstance().disableExtension((String) REQUEST_META.get().get("requestURI"))) {
return response;
}
if (ConfigMatcher.getInstance().getBlackUrl(REQUEST_META.get())) {
return response;
}
if (cloneResponseMethod == null) {
loadCloneResponseMethod();
}
Expand All @@ -114,6 +121,8 @@ public static Object cloneResponse(Object response, boolean isJakarta) {
return response;
} catch (InvocationTargetException e) {
return response;
} finally {
REQUEST_META.remove();
}
}

Expand Down Expand Up @@ -152,22 +161,22 @@ public static Map<String, Object> getResponseMeta(Object response) {
public static void solveHttp(MethodEvent event)
throws InvocationTargetException, IllegalAccessException, NoSuchMethodException {
DongTaiLog.debug(EngineManager.SCOPE_TRACKER.get().toString());
Map<String, Object> requestMeta = getRequestMeta(event.argumentArray[0]);
Boolean isReplay = (Boolean) requestMeta.get("replay-request");
REQUEST_META.set(getRequestMeta(event.argumentArray[0]));
Boolean isReplay = (Boolean) REQUEST_META.get().get("replay-request");
if (isReplay){
EngineManager.ENTER_REPLAY_ENTRYPOINT.enterEntry();
}
// todo Consider increasing the capture of html request responses
if (ConfigMatcher.getInstance().disableExtension((String) requestMeta.get("requestURI"))) {
if (ConfigMatcher.getInstance().disableExtension((String) REQUEST_META.get().get("requestURI"))) {
return;
}
if (ConfigMatcher.getInstance().getBlackUrl(requestMeta)) {
if (ConfigMatcher.getInstance().getBlackUrl(REQUEST_META.get())) {
return;
}

// todo: add custom header escape
EngineManager.enterHttpEntry(requestMeta);
DongTaiLog.debug("HTTP Request:{} {} from: {}", requestMeta.get("method"), requestMeta.get("requestURI"),
EngineManager.enterHttpEntry(REQUEST_META.get());
DongTaiLog.debug("HTTP Request:{} {} from: {}", REQUEST_META.get().get("method"), REQUEST_META.get().get("requestURI"),
event.signature);
}

Expand Down
2 changes: 1 addition & 1 deletion dongtai-core/src/main/resources/com.secnium.iast.resources/blackext.txt
View file
Original file line number Diff line number Diff line change
@@ -1 +1 @@
.js,.css,.htm,.html,.jpg,.png,.gif,.woff,.woff2,.ico,.maps,.xml
.js,.css,.htm,.html,.jpg,.png,.gif,.woff,.woff2,.ico,.maps,.xml,.map

0 comments on commit 8417960

Please sign in to comment.