Merge pull request #245 from samwafgo/feat_modify_task

feat:task check exist logic

service/waf_service/waf_task_service.go

@@ -81,7 +81,7 @@ func (receiver *WafTaskService) CheckIsExistApi(req request.WafTaskAddReq) int {
 	return int(total)
 }
 
-func (receiver *WafTaskService) CheckIsExist(taskName string) int {
+func (receiver *WafTaskService) CheckIsExist(taskMethod string) int {
 	var total int64 = 0
 	/*where条件*/
 	var whereField = ""
@@ -89,17 +89,17 @@ func (receiver *WafTaskService) CheckIsExist(taskName string) int {
 	//where字段
 	whereField = ""
 
-	if len(taskName) > 0 {
+	if len(taskMethod) > 0 {
 		if len(whereField) > 0 {
 			whereField = whereField + " and "
 		}
-		whereField = whereField + " task_name=? "
+		whereField = whereField + " task_method=? "
 	}
 
 	//where字段赋值
-	if len(taskName) > 0 {
+	if len(taskMethod) > 0 {
 		if len(whereField) > 0 {
-			whereValues = append(whereValues, taskName)
+			whereValues = append(whereValues, taskMethod)
 		}
 	}
 
@@ -188,3 +188,12 @@ func (receiver *WafTaskService) DelApi(req request.WafTaskDelReq) error {
 	err = global.GWAF_LOCAL_DB.Where("id = ?", req.Id).Delete(model.Task{}).Error
 	return err
 }
+func (receiver *WafTaskService) DelByMethod(taskMethod string) error {
+	var bean model.Task
+	err := global.GWAF_LOCAL_DB.Where("task_method = ?", taskMethod).First(&bean).Error
+	if err != nil {
+		return err
+	}
+	err = global.GWAF_LOCAL_DB.Where("task_method = ?", taskMethod).Delete(model.Task{}).Error
+	return err
+}

waftask/entrance.go

@@ -128,8 +128,12 @@ func InitTaskDb() []model.Task {
 	return list
 }
 func syncTaskToDb(task model.Task) {
-	cnt := wafTaskService.CheckIsExist(task.TaskName)
+	cnt := wafTaskService.CheckIsExist(task.TaskMethod)
 	if cnt == 0 {
 		wafTaskService.Add(task)
+	} else if cnt > 1 {
+		//如果重复添加多个method的情况
+		wafTaskService.DelByMethod(task.TaskMethod)
+		wafTaskService.Add(task)
 	}
 }