修复栈溢出

Author: keven1z

dhook.sqlite

@@ -66,6 +66,7 @@ public int update(@RequestParam(value = "id") String id, String name) {
      */
     @GetMapping("/del")
     public int del(String agentId) {
+        classMapService.delete(agentId);
         return agentService.delete(agentId);
     }
 

src/main/java/com/keven1z/controller/AgentController.java

@@ -130,8 +130,8 @@ public ResponseEntity<Object> export(@RequestParam(value = "id") String id, Stri
     }
 
     @PostMapping("/detail/add")
-    public int addHookDetail(@RequestBody HookDetailEntity hookDetailEntity) {
-        return hookService.addHookDetail(hookDetailEntity);
+    public int addHookDetail(@RequestBody List<HookDetailEntity> hookDetailEntity) {
+        return hookService.addHookDetails(hookDetailEntity);
     }
 
     @GetMapping("/detail/find")
@@ -177,4 +177,8 @@ private static boolean inject(VirtualMachineDescriptor vmd,String agentPath) thr
         }
         return false;
     }
+
+//    public static void main(String[] args) throws Exception {
+//        inject("cn.com.tcsec.goat.taint.TaintApplication","C:\\Users\\fbi\\Desktop\\gitProject\\DHookAgent\\target\\dHook.jar");
+//    }
 }

src/main/java/com/keven1z/controller/HookController.java

@@ -1,5 +1,6 @@
 package com.keven1z.dao;
 
+import com.keven1z.entity.ClassMapEntity;
 import com.keven1z.entity.HookDetailEntity;
 import com.keven1z.entity.HookEntity;
 import org.apache.ibatis.annotations.*;
@@ -93,6 +94,16 @@ public interface IHookDao {
     @Options(useGeneratedKeys=true, keyProperty="id")
     int addHookDetail(HookDetailEntity hookDetail);
 
+    @Insert("<script>"  +
+            "INSERT or ignore INTO hook_detail(hook_id, param, returnObject, thisObject,date,stacks) VALUES" +
+            "<foreach collection=\"hookDetailEntities\" item=\"item1\" index=\"index\"  separator=\",\">" +
+            "(#{item1.hookId},#{item1.param},#{item1.returnObject},#{item1.thisObject},#{item1.date},#{item1.stacks})" +
+            "</foreach>" +
+            "</script>")
+    @Result(column = "hook_id",property = "hookId")
+    @Options(useGeneratedKeys=true, keyProperty="id")
+    int addDetails(@Param(value="hookDetailEntities") List<HookDetailEntity> hookDetailEntities);
+
     @Select("SELECT * FROM hook_detail WHERE id = #{id} order by date desc ")
     @Results(value={
             @Result(id = true,column = "id",property = "id"),

src/main/java/com/keven1z/dao/IHookDao.java

@@ -70,7 +70,11 @@ protected void channelRead0(ChannelHandlerContext ctx, CustomProtocol customProt
             NettySocketHolder.put(customProtocol.getId(), (NioSocketChannel) ctx.channel());
             String body = customProtocol.getBody();
             AgentEntity agentEntity = GsonUtil.toBean(body, AgentEntity.class);
-            AgentUtil.register(agentEntity);
+            try {
+                AgentUtil.register(agentEntity);
+            } catch (Exception e) {
+                LOGGER.error("Register agent failed!reason:" + e.getMessage());
+            }
         }
 
         CustomProtocol protocol = HeartbeatInitializer.HeartQueue.poll();

src/main/java/com/keven1z/netty/HeartBeatSimpleHandle.java

@@ -30,4 +30,5 @@ public interface IHookService {
     HookDetailEntity findHookDetailById(int id);
     List<HookDetailEntity> findHookDetails();
     int addHookDetail(HookDetailEntity hookDetail);
+    int addHookDetails(List<HookDetailEntity> hookDetails);
 }

src/main/java/com/keven1z/service/IHookService.java

@@ -145,6 +145,10 @@ public List<HookDetailEntity> findHookDetails() {
     public int addHookDetail(HookDetailEntity hookDetail) {
         return hookDao.addHookDetail(hookDetail);
     }
+    @Override
+    public int addHookDetails(List<HookDetailEntity> hookDetails) {
+        return hookDao.addDetails(hookDetails);
+    }
 
     public boolean isInField(int fieldId,List<FieldEntity> fieldEntityList){
         for (FieldEntity fieldEntity:fieldEntityList){

src/main/java/com/keven1z/service/impl/HookServiceImpl.java

@@ -2,6 +2,7 @@
 
 import com.keven1z.entity.AgentEntity;
 import com.keven1z.service.IAgentService;
+import com.keven1z.service.IClassMapService;
 import org.springframework.stereotype.Component;
 
 import javax.annotation.PostConstruct;
@@ -15,6 +16,8 @@
 public class AgentUtil {
     @Resource
     private IAgentService agentService;
+    @Resource
+    private IClassMapService classMapService;
 
     public static AgentUtil agentUtil;
 
@@ -28,5 +31,6 @@ public static void update(String id, int state, String time) {
     }
     public static void register(AgentEntity agentEntity) {
         agentUtil.agentService.register(agentEntity);
+        agentUtil.classMapService.delete(agentEntity.getId());
     }
 }