From 415b76cb4e1037ff46e42ad890da4082803beaa4 Mon Sep 17 00:00:00 2001 From: samwaf Date: Mon, 26 May 2025 15:31:48 +0800 Subject: [PATCH] fix:batch insert ip slow #329 --- waftask/batch/batch_processor.go | 237 ++++++++++++++++++++++++++++ waftask/batch/ip_allow_processor.go | 178 +++++++++++++++++++++ waftask/batch/ip_deny_processor.go | 178 +++++++++++++++++++++ waftask/batch/item_extractor.go | 68 ++++++++ waftask/task_batch.go | 186 ++-------------------- 5 files changed, 673 insertions(+), 174 deletions(-) create mode 100644 waftask/batch/batch_processor.go create mode 100644 waftask/batch/ip_allow_processor.go create mode 100644 waftask/batch/ip_deny_processor.go create mode 100644 waftask/batch/item_extractor.go diff --git a/waftask/batch/batch_processor.go b/waftask/batch/batch_processor.go new file mode 100644 index 0000000..dff1673 --- /dev/null +++ b/waftask/batch/batch_processor.go @@ -0,0 +1,237 @@ +package batch + +import ( + "SamWaf/common/zlog" + "SamWaf/global" + "SamWaf/model" + "SamWaf/utils" + "bufio" + "fmt" + "io" + "net/http" + "os" + "regexp" + "strings" + "sync/atomic" +) + +// BatchProcessor 批量处理器接口 +type BatchProcessor interface { + // ProcessItem 处理单个项目 + ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool + // GetExistingItems 获取已存在的项目 + GetExistingItems(items []string, task model.BatchTask) map[string]interface{} + // NotifyEngine 通知引擎更新 + NotifyEngine(task model.BatchTask) +} + +// BatchProcessorConfig 批量处理器配置 +type BatchProcessorConfig struct { + BatchSize int // 批处理大小 + LogPrefix string // 日志前缀 +} + +// BatchProgress 批处理进度 +type BatchProgress struct { + TotalItems int32 // 总项目数 + ProcessedItems int32 // 已处理项目数 + InsertedItems int32 // 已插入项目数 + UpdatedItems int32 // 已更新项目数 +} + +// AddProcessed 增加已处理数量 +func (p *BatchProgress) AddProcessed(count int) { + atomic.AddInt32(&p.ProcessedItems, int32(count)) +} + +// AddInserted 增加已插入数量 +func (p *BatchProgress) AddInserted(count int) { + atomic.AddInt32(&p.InsertedItems, int32(count)) +} + +// AddUpdated 增加已更新数量 +func (p *BatchProgress) AddUpdated(count int) { + atomic.AddInt32(&p.UpdatedItems, int32(count)) +} + +// GetProgress 获取进度百分比 +func (p *BatchProgress) GetProgress() float64 { + if p.TotalItems == 0 { + return 100.0 + } + return float64(p.ProcessedItems) / float64(p.TotalItems) * 100.0 +} + +// ProcessBatchTask 通用批量处理函数 +func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config BatchProcessorConfig) { + innerLogName := config.LogPrefix + + contentReader, err := openSource(task) + if err != nil { + zlog.Error(innerLogName, err.Error()) + return + } + defer contentReader.Close() + + // 判断数据库是否已经关闭 + if global.GWAF_LOCAL_DB == nil { + zlog.Error(innerLogName, "数据库已经关闭批量处理终止") + return + } + + // 首先计算总行数，用于进度显示 + totalLines, validLines, err := countLines(task) + if err != nil { + zlog.Error(innerLogName, fmt.Sprintf("计算总行数失败: %s", err.Error())) + // 继续执行，但无法显示准确进度 + } + + // 创建进度跟踪器 + progress := &BatchProgress{ + TotalItems: int32(validLines), + } + + zlog.Info(innerLogName, fmt.Sprintf("开始批量处理，总行数: %d，有效行数: %d", totalLines, validLines)) + + // 收集有效的项目 + validItems := make([]string, 0, config.BatchSize) + scanner := bufio.NewScanner(contentReader) + hasAffectInfo := false + + // 按批次处理 + batchCount := 0 + for scanner.Scan() { + line := scanner.Text() + line = strings.TrimSpace(line) + if line == "" { + continue // 跳过空行 + } + + line = extractIPFromLine(line) // 可以根据需要替换为更通用的提取函数 + validRet, _ := utils.IsValidIPOrNetwork(line) // 可以根据需要替换为更通用的验证函数 + if !validRet { + continue + } + + validItems = append(validItems, line) + batchCount++ + + // 当收集到一批或者是最后一批时，进行批量处理 + if batchCount >= config.BatchSize { + // 处理当前批次 + if processor.ProcessBatch(validItems, task, progress) { + hasAffectInfo = true + } + + // 记录进度 + progress.AddProcessed(len(validItems)) + zlog.Info(innerLogName, fmt.Sprintf("进度: %.2f%% (%d/%d)，已插入: %d，已更新: %d", + progress.GetProgress(), progress.ProcessedItems, progress.TotalItems, + progress.InsertedItems, progress.UpdatedItems)) + + // 重置批次 + validItems = make([]string, 0, config.BatchSize) + batchCount = 0 + } + } + + // 处理最后一批不足BatchSize的项目 + if len(validItems) > 0 { + if processor.ProcessBatch(validItems, task, progress) { + hasAffectInfo = true + } + + // 记录进度 + progress.AddProcessed(len(validItems)) + zlog.Info(innerLogName, fmt.Sprintf("进度: %.2f%% (%d/%d)，已插入: %d，已更新: %d", + progress.GetProgress(), progress.ProcessedItems, progress.TotalItems, + progress.InsertedItems, progress.UpdatedItems)) + } + + if hasAffectInfo { + // 通知引擎进行实时生效 + processor.NotifyEngine(task) + } + + // 输出最终处理结果 + zlog.Info(innerLogName, fmt.Sprintf("批量处理完成，总计处理: %d，插入: %d，更新: %d", + progress.ProcessedItems, progress.InsertedItems, progress.UpdatedItems)) + + // 检查扫描错误 + if err := scanner.Err(); err != nil { + zlog.Error(innerLogName, fmt.Sprintf("扫描文件时发生错误: %s", err.Error())) + } +} + +// countLines 计算文件总行数和有效行数 +func countLines(task model.BatchTask) (int, int, error) { + contentReader, err := openSource(task) + if err != nil { + return 0, 0, err + } + defer contentReader.Close() + + scanner := bufio.NewScanner(contentReader) + totalLines := 0 + validLines := 0 + + for scanner.Scan() { + line := scanner.Text() + totalLines++ + + line = strings.TrimSpace(line) + if line == "" { + continue // 跳过空行 + } + + line = extractIPFromLine(line) + validRet, _ := utils.IsValidIPOrNetwork(line) + if validRet { + validLines++ + } + } + + return totalLines, validLines, scanner.Err() +} + +// openSource 打开本地或远程数据源 +func openSource(task model.BatchTask) (io.ReadCloser, error) { + if task.BatchSourceType == "local" { + f, err := os.Open(task.BatchSource) + if err != nil { + return nil, fmt.Errorf("failed to open local file: %v", err) + } + return f, nil + } + // remote + resp, err := http.Get(task.BatchSource) + if err != nil { + return nil, fmt.Errorf("failed to fetch remote data: %v", err) + } + if resp.StatusCode != http.StatusOK { + resp.Body.Close() + return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode) + } + return resp.Body, nil +} + +// extractIPFromLine 使用正则表达式从行中提取IP地址或网段 +func extractIPFromLine(line string) string { + // 匹配IPv4地址或IPv4网段 + ipv4Regex := regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}(?:/\d{1,2})?\b`) + + // 匹配IPv6地址或IPv6网段 (简化版本) + ipv6Regex := regexp.MustCompile(`\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}(?:/\d{1,3})?\b`) + + // 先尝试匹配IPv4 + if match := ipv4Regex.FindString(line); match != "" { + return match + } + + // 再尝试匹配IPv6 + if match := ipv6Regex.FindString(line); match != "" { + return match + } + + return line // 如果没有匹配到，返回原始行 +} diff --git a/waftask/batch/ip_allow_processor.go b/waftask/batch/ip_allow_processor.go new file mode 100644 index 0000000..cd3f005 --- /dev/null +++ b/waftask/batch/ip_allow_processor.go @@ -0,0 +1,178 @@ +package batch + +import ( + "SamWaf/common/uuid" + "SamWaf/common/zlog" + "SamWaf/customtype" + "SamWaf/enums" + "SamWaf/global" + "SamWaf/model" + "SamWaf/model/baseorm" + "SamWaf/model/spec" + "fmt" + "time" +) + +// IPAllowProcessor 白名单IP处理器 +type IPAllowProcessor struct{} + +// ProcessBatch 处理一批IP +func (p *IPAllowProcessor) ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool { + if len(items) == 0 { + return false + } + + logName := "BatchTask-IPAllowBatch" + zlog.Info(logName, fmt.Sprintf("处理白名单批次，包含 %d 个IP", len(items))) + + // 获取已存在的记录 + existMap := p.GetExistingItems(items, task) + + // 根据执行方法处理 + if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND { + return p.processAppendBatch(items, existMap, task, logName, progress) + } else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE { + return p.processOverwriteBatch(items, existMap, task, logName, progress) + } + + return false +} + +// GetExistingItems 获取已存在的IP记录 +func (p *IPAllowProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} { + existMap := make(map[string]interface{}) + var existIPs []model.IPAllowList + + // 使用IN查询一次性获取所有已存在的记录 + global.GWAF_LOCAL_DB.Where("host_code = ? AND ip IN (?)", task.BatchHostCode, items).Find(&existIPs) + + // 将已存在的IP放入map中，方便快速查找 + for _, ip := range existIPs { + existMap[ip.Ip] = ip + } + + return existMap +} + +// NotifyEngine 通知引擎更新 +func (p *IPAllowProcessor) NotifyEngine(task model.BatchTask) { + var ipWhites []model.IPAllowList + global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipWhites) + var chanInfo = spec.ChanCommonHost{ + HostCode: task.BatchHostCode, + Type: enums.ChanTypeAllowIP, + Content: ipWhites, + } + global.GWAF_CHAN_MSG <- chanInfo +} + +// processAppendBatch 处理追加模式的批次 +func (p *IPAllowProcessor) processAppendBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool { + // 收集需要插入的记录 + var toInsert []model.IPAllowList + for _, ip := range items { + // 如果IP不存在，则添加到待插入列表 + if _, exists := existMap[ip]; !exists { + toInsert = append(toInsert, model.IPAllowList{ + BaseOrm: baseorm.BaseOrm{ + Id: uuid.GenUUID(), + USER_CODE: global.GWAF_USER_CODE, + Tenant_ID: global.GWAF_TENANT_ID, + CREATE_TIME: customtype.JsonTime(time.Now()), + UPDATE_TIME: customtype.JsonTime(time.Now()), + }, + HostCode: task.BatchHostCode, + Ip: ip, + Remarks: time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id, + }) + } + } + + // 批量插入新记录 + if len(toInsert) > 0 { + tx := global.GWAF_LOCAL_DB.Begin() + if err := tx.Create(&toInsert).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量插入白名单IP失败: "+err.Error()) + return false + } + tx.Commit() + + zlog.Info(logName, fmt.Sprintf("成功插入 %d 条白名单IP记录", len(toInsert))) + // 更新进度统计 + progress.AddInserted(len(toInsert)) + return len(toInsert) > 0 + } + + return false +} + +// processOverwriteBatch 处理覆写模式的批次 +func (p *IPAllowProcessor) processOverwriteBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool { + // 收集需要插入和更新的记录 + var toInsert []model.IPAllowList + var toUpdate []model.IPAllowList + + for _, ip := range items { + if existIP, exists := existMap[ip]; exists { + // 已存在，需要更新 + ipRecord := existIP.(model.IPAllowList) + ipRecord.Remarks = time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id + ipRecord.UPDATE_TIME = customtype.JsonTime(time.Now()) + toUpdate = append(toUpdate, ipRecord) + } else { + // 不存在，需要插入 + toInsert = append(toInsert, model.IPAllowList{ + BaseOrm: baseorm.BaseOrm{ + Id: uuid.GenUUID(), + USER_CODE: global.GWAF_USER_CODE, + Tenant_ID: global.GWAF_TENANT_ID, + CREATE_TIME: customtype.JsonTime(time.Now()), + UPDATE_TIME: customtype.JsonTime(time.Now()), + }, + HostCode: task.BatchHostCode, + Ip: ip, + Remarks: time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id, + }) + } + } + + // 开始事务 + tx := global.GWAF_LOCAL_DB.Begin() + hasChanges := false + + // 批量插入新记录 + if len(toInsert) > 0 { + if err := tx.Create(&toInsert).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量插入白名单IP失败: "+err.Error()) + return false + } + zlog.Info(logName, fmt.Sprintf("成功插入 %d 条白名单IP记录", len(toInsert))) + // 更新进度统计 + progress.AddInserted(len(toInsert)) + hasChanges = true + } + + // 批量更新已存在的记录 + if len(toUpdate) > 0 { + for _, record := range toUpdate { + if err := tx.Model(&model.IPAllowList{}).Where("id = ?", record.Id).Updates(map[string]interface{}{ + "Remarks": record.Remarks, + "UPDATE_TIME": record.UPDATE_TIME, + }).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量更新白名单IP失败: "+err.Error()) + return false + } + } + zlog.Info(logName, fmt.Sprintf("成功更新 %d 条白名单IP记录", len(toUpdate))) + // 更新进度统计 + progress.AddUpdated(len(toUpdate)) + hasChanges = true + } + + // 提交事务 + tx.Commit() + return hasChanges +} diff --git a/waftask/batch/ip_deny_processor.go b/waftask/batch/ip_deny_processor.go new file mode 100644 index 0000000..49aba7c --- /dev/null +++ b/waftask/batch/ip_deny_processor.go @@ -0,0 +1,178 @@ +package batch + +import ( + "SamWaf/common/uuid" + "SamWaf/common/zlog" + "SamWaf/customtype" + "SamWaf/enums" + "SamWaf/global" + "SamWaf/model" + "SamWaf/model/baseorm" + "SamWaf/model/spec" + "fmt" + "time" +) + +// IPDenyProcessor 黑名单IP处理器 +type IPDenyProcessor struct{} + +// ProcessBatch 处理一批IP +func (p *IPDenyProcessor) ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool { + if len(items) == 0 { + return false + } + + logName := "BatchTask-IPDenyBatch" + zlog.Info(logName, fmt.Sprintf("处理黑名单批次，包含 %d 个IP", len(items))) + + // 获取已存在的记录 + existMap := p.GetExistingItems(items, task) + + // 根据执行方法处理 + if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND { + return p.processAppendBatch(items, existMap, task, logName, progress) + } else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE { + return p.processOverwriteBatch(items, existMap, task, logName, progress) + } + + return false +} + +// GetExistingItems 获取已存在的IP记录 +func (p *IPDenyProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} { + existMap := make(map[string]interface{}) + var existIPs []model.IPBlockList + + // 使用IN查询一次性获取所有已存在的记录 + global.GWAF_LOCAL_DB.Where("host_code = ? AND ip IN (?)", task.BatchHostCode, items).Find(&existIPs) + + // 将已存在的IP放入map中，方便快速查找 + for _, ip := range existIPs { + existMap[ip.Ip] = ip + } + + return existMap +} + +// NotifyEngine 通知引擎更新 +func (p *IPDenyProcessor) NotifyEngine(task model.BatchTask) { + var ipBlocks []model.IPBlockList + global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipBlocks) + var chanInfo = spec.ChanCommonHost{ + HostCode: task.BatchHostCode, + Type: enums.ChanTypeBlockIP, + Content: ipBlocks, + } + global.GWAF_CHAN_MSG <- chanInfo +} + +// processAppendBatch 处理追加模式的批次 +func (p *IPDenyProcessor) processAppendBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool { + // 收集需要插入的记录 + var toInsert []model.IPBlockList + for _, ip := range items { + // 如果IP不存在，则添加到待插入列表 + if _, exists := existMap[ip]; !exists { + toInsert = append(toInsert, model.IPBlockList{ + BaseOrm: baseorm.BaseOrm{ + Id: uuid.GenUUID(), + USER_CODE: global.GWAF_USER_CODE, + Tenant_ID: global.GWAF_TENANT_ID, + CREATE_TIME: customtype.JsonTime(time.Now()), + UPDATE_TIME: customtype.JsonTime(time.Now()), + }, + HostCode: task.BatchHostCode, + Ip: ip, + Remarks: time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id, + }) + } + } + + // 批量插入新记录 + if len(toInsert) > 0 { + tx := global.GWAF_LOCAL_DB.Begin() + if err := tx.Create(&toInsert).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量插入黑名单IP失败: "+err.Error()) + return false + } + tx.Commit() + + zlog.Info(logName, fmt.Sprintf("成功插入 %d 条黑名单IP记录", len(toInsert))) + // 更新进度统计 + progress.AddInserted(len(toInsert)) + return len(toInsert) > 0 + } + + return false +} + +// processOverwriteBatch 处理覆写模式的批次 +func (p *IPDenyProcessor) processOverwriteBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool { + // 收集需要插入和更新的记录 + var toInsert []model.IPBlockList + var toUpdate []model.IPBlockList + + for _, ip := range items { + if existIP, exists := existMap[ip]; exists { + // 已存在，需要更新 + ipRecord := existIP.(model.IPBlockList) + ipRecord.Remarks = time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id + ipRecord.UPDATE_TIME = customtype.JsonTime(time.Now()) + toUpdate = append(toUpdate, ipRecord) + } else { + // 不存在，需要插入 + toInsert = append(toInsert, model.IPBlockList{ + BaseOrm: baseorm.BaseOrm{ + Id: uuid.GenUUID(), + USER_CODE: global.GWAF_USER_CODE, + Tenant_ID: global.GWAF_TENANT_ID, + CREATE_TIME: customtype.JsonTime(time.Now()), + UPDATE_TIME: customtype.JsonTime(time.Now()), + }, + HostCode: task.BatchHostCode, + Ip: ip, + Remarks: time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id, + }) + } + } + + // 开始事务 + tx := global.GWAF_LOCAL_DB.Begin() + hasChanges := false + + // 批量插入新记录 + if len(toInsert) > 0 { + if err := tx.Create(&toInsert).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量插入黑名单IP失败: "+err.Error()) + return false + } + zlog.Info(logName, fmt.Sprintf("成功插入 %d 条黑名单IP记录", len(toInsert))) + // 更新进度统计 + progress.AddInserted(len(toInsert)) + hasChanges = true + } + + // 批量更新已存在的记录 + if len(toUpdate) > 0 { + for _, record := range toUpdate { + if err := tx.Model(&model.IPBlockList{}).Where("id = ?", record.Id).Updates(map[string]interface{}{ + "Remarks": record.Remarks, + "UPDATE_TIME": record.UPDATE_TIME, + }).Error; err != nil { + tx.Rollback() + zlog.Error(logName, "批量更新黑名单IP失败: "+err.Error()) + return false + } + } + zlog.Info(logName, fmt.Sprintf("成功更新 %d 条黑名单IP记录", len(toUpdate))) + // 更新进度统计 + progress.AddUpdated(len(toUpdate)) + hasChanges = true + } + + // 提交事务 + tx.Commit() + return hasChanges +} diff --git a/waftask/batch/item_extractor.go b/waftask/batch/item_extractor.go new file mode 100644 index 0000000..719ff3e --- /dev/null +++ b/waftask/batch/item_extractor.go @@ -0,0 +1,68 @@ +package batch + +import ( + "SamWaf/utils" + "regexp" + "strings" +) + +// ItemExtractor 项目提取器接口 +type ItemExtractor interface { + // ExtractItem 从行中提取项目 + ExtractItem(line string) string + // ValidateItem 验证项目是否有效 + ValidateItem(item string) bool +} + +// IPExtractor IP提取器 +type IPExtractor struct{} + +// ExtractItem 从行中提取IP地址 +func (e *IPExtractor) ExtractItem(line string) string { + // 匹配IPv4地址或IPv4网段 + ipv4Regex := regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}(?:/\d{1,2})?\b`) + + // 匹配IPv6地址或IPv6网段 (简化版本) + ipv6Regex := regexp.MustCompile(`\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}(?:/\d{1,3})?\b`) + + // 先尝试匹配IPv4 + if match := ipv4Regex.FindString(line); match != "" { + return match + } + + // 再尝试匹配IPv6 + if match := ipv6Regex.FindString(line); match != "" { + return match + } + + return line // 如果没有匹配到，返回原始行 +} + +// ValidateItem 验证IP地址是否有效 +func (e *IPExtractor) ValidateItem(item string) bool { + validRet, _ := utils.IsValidIPOrNetwork(item) + return validRet +} + +// DefaultExtractor 默认提取器，不做特殊处理 +type DefaultExtractor struct{} + +// ExtractItem 默认提取，只做简单的空格处理 +func (e *DefaultExtractor) ExtractItem(line string) string { + return strings.TrimSpace(line) +} + +// ValidateItem 默认验证，非空即有效 +func (e *DefaultExtractor) ValidateItem(item string) bool { + return item != "" +} + +// GetExtractor 根据批量任务类型获取合适的提取器 +func GetExtractor(batchType string) ItemExtractor { + switch batchType { + case "ipallow", "ipdeny": + return &IPExtractor{} + default: + return &DefaultExtractor{} + } +} diff --git a/waftask/task_batch.go b/waftask/task_batch.go index eed8867..abd2945 100644 --- a/waftask/task_batch.go +++ b/waftask/task_batch.go @@ -3,25 +3,13 @@ package waftask import ( "SamWaf/common/zlog" "SamWaf/enums" - "SamWaf/global" "SamWaf/model" - "SamWaf/model/request" - "SamWaf/model/spec" "SamWaf/service/waf_service" - "SamWaf/utils" - "errors" - "fmt" - "gorm.io/gorm" - "io/ioutil" - "net/http" - "strings" - "time" + "SamWaf/waftask/batch" ) var ( - wafBatchTaskService = waf_service.WafBatchServiceApp - wafIPAllowService = waf_service.WafWhiteIpServiceApp - wafBlockAllowService = waf_service.WafBlockIpServiceApp + wafBatchTaskService = waf_service.WafBatchServiceApp ) /* @@ -56,172 +44,22 @@ func BatchTask() { } } -// handleLocalSource 从本地路径读取数据 -func handleLocalSource(task model.BatchTask) (string, error) { - data, err := ioutil.ReadFile(task.BatchSource) - if err != nil { - return "", fmt.Errorf("failed to read local file: %v", err) - } - return string(data), nil -} - -// handleRemoteSource 从远程 URL 获取数据 -func handleRemoteSource(task model.BatchTask) (string, error) { - resp, err := http.Get(task.BatchSource) - if err != nil { - return "", fmt.Errorf("failed to fetch remote data: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("unexpected status code: %d", resp.StatusCode) - } - - data, err := ioutil.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("failed to read remote response body: %v", err) - } - return string(data), nil -} - // IPAllowBatch 白名单IP批量处理 func IPAllowBatch(task model.BatchTask) { - innerLogName := "BatchTask-IPAllowBatch" - //1.取数据 - //2.处理数据 - content := "" - if task.BatchSourceType == "local" { - cnt, err := handleLocalSource(task) - if err != nil { - zlog.Error(innerLogName, err.Error()) - return - } - content = cnt - } else if task.BatchSourceType == "remote" { - cnt, err := handleRemoteSource(task) - if err != nil { - zlog.Error(innerLogName, err.Error()) - return - } - content = cnt + processor := &batch.IPAllowProcessor{} + config := batch.BatchProcessorConfig{ + BatchSize: 1000, + LogPrefix: "BatchTask-IPAllowBatch", } - if content == "" { - zlog.Error(innerLogName, task.BatchTaskName+"没有数据需要处理") - return - } - hasAffectInfo := false - lines := strings.Split(content, "\n") - // 遍历每一行数据进行处理 - for _, line := range lines { - if line == "" { - continue // 跳过空行 - } - line = strings.TrimSpace(line) - validRet, _ := utils.IsValidIPOrNetwork(line) - if !validRet { - continue - } - if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND { - //添加 - err := wafIPAllowService.CheckIsExistApi(request.WafAllowIpAddReq{task.BatchHostCode, line, ""}) - if err != nil && errors.Is(err, gorm.ErrRecordNotFound) { - wafIPAllowService.AddApi(request.WafAllowIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id}) - hasAffectInfo = true - } - } else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE { - //覆写 - bean := wafIPAllowService.GetDetailByIPApi(line, task.BatchHostCode) - if bean.HostCode == "" { - wafIPAllowService.AddApi(request.WafAllowIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id}) - hasAffectInfo = true - } else { - wafIPAllowService.ModifyApi(request.WafAllowIpEditReq{bean.Id, task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id}) - hasAffectInfo = true - } - - } - } - if hasAffectInfo { - //发送通知到引擎进行实时生效 - var ipWhites []model.IPAllowList - global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipWhites) - var chanInfo = spec.ChanCommonHost{ - HostCode: task.BatchHostCode, - Type: enums.ChanTypeAllowIP, - Content: ipWhites, - } - global.GWAF_CHAN_MSG <- chanInfo - } - + batch.ProcessBatchTask(task, processor, config) } // IPDenyBatch 黑名单IP批量处理 func IPDenyBatch(task model.BatchTask) { - innerLogName := "BatchTask-IPDenyBatch" - //1.取数据 - //2.处理数据 - content := "" - if task.BatchSourceType == "local" { - cnt, err := handleLocalSource(task) - if err != nil { - zlog.Error(innerLogName, err.Error()) - return - } - content = cnt - } else if task.BatchSourceType == "remote" { - cnt, err := handleRemoteSource(task) - if err != nil { - zlog.Error(innerLogName, err.Error()) - return - } - content = cnt + processor := &batch.IPDenyProcessor{} + config := batch.BatchProcessorConfig{ + BatchSize: 1000, + LogPrefix: "BatchTask-IPDenyBatch", } - if content == "" { - zlog.Error(innerLogName, task.BatchTaskName+"没有数据需要处理") - return - } - hasAffectInfo := false - lines := strings.Split(content, "\n") - // 遍历每一行数据进行处理 - for _, line := range lines { - if line == "" { - continue // 跳过空行 - } - line = strings.TrimSpace(line) - validRet, _ := utils.IsValidIPOrNetwork(line) - if !validRet { - continue - } - if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND { - //添加 - err := wafBlockAllowService.CheckIsExistApi(request.WafBlockIpAddReq{task.BatchHostCode, line, ""}) - if err != nil && errors.Is(err, gorm.ErrRecordNotFound) { - wafBlockAllowService.AddApi(request.WafBlockIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id}) - hasAffectInfo = true - } - } else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE { - //覆写 - bean := wafBlockAllowService.GetDetailByIPApi(line, task.BatchHostCode) - if bean.HostCode == "" { - wafBlockAllowService.AddApi(request.WafBlockIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id}) - hasAffectInfo = true - } else { - wafBlockAllowService.ModifyApi(request.WafBlockIpEditReq{bean.Id, task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id}) - hasAffectInfo = true - } - - } - } - if hasAffectInfo { - //发送通知到引擎进行实时生效 - var ipBlocks []model.IPBlockList - global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipBlocks) - var chanInfo = spec.ChanCommonHost{ - HostCode: task.BatchHostCode, - Type: enums.ChanTypeBlockIP, - Content: ipBlocks, - } - global.GWAF_CHAN_MSG <- chanInfo - } - + batch.ProcessBatchTask(task, processor, config) }