mirror of
https://gitee.com/samwaf/SamWaf.git
synced 2026-10-31 10:11:42 +08:00
396 lines
13 KiB
Go
396 lines
13 KiB
Go
package batch
|
||
|
||
import (
|
||
"SamWaf/common/zlog"
|
||
"SamWaf/enums"
|
||
"SamWaf/global"
|
||
"SamWaf/model"
|
||
"SamWaf/utils"
|
||
"bufio"
|
||
"bytes"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"regexp"
|
||
"strings"
|
||
"sync/atomic"
|
||
"time"
|
||
)
|
||
|
||
// BatchProcessor 批量处理器接口
|
||
type BatchProcessor interface {
|
||
// ProcessItem 处理单个项目
|
||
ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool
|
||
// GetExistingItems 获取已存在的项目
|
||
GetExistingItems(items []string, task model.BatchTask, config interface{}) map[string]interface{}
|
||
// NotifyEngine 通知引擎更新
|
||
NotifyEngine(task model.BatchTask)
|
||
}
|
||
|
||
// BatchFinalizer 可选接口：所有批次都跑完之后回调一次。
|
||
//
|
||
// 给「必须看完整份源才能决定」的收尾用（例如覆盖模式下删除源里已经不存在的旧条目）。
|
||
// sourceComplete 表示这份源是否被完整读完（扫描无错误）——半截数据绝不能拿去做删除，
|
||
// 否则一次网络抖动就能把组清空：被白名单引用时是全站失去豁免，被黑名单引用时是瞬间放行。
|
||
// 返回值语义与 ProcessBatch 一致：true 表示数据有变化，需要通知引擎。
|
||
type BatchFinalizer interface {
|
||
Finalize(task model.BatchTask, progress *BatchProgress, sourceComplete bool) bool
|
||
}
|
||
|
||
// BatchProcessorConfig 批量处理器配置
|
||
type BatchProcessorConfig struct {
|
||
BatchSize int // 批处理大小
|
||
LogPrefix string // 日志前缀
|
||
// MaxLines 本次导入允许的最大行数，<=0 时取 defaultMaxSourceLines。
|
||
// 敏感词这类"导进来立刻进检测引擎、每个请求都要逐条匹配"的类型需要单独收紧。
|
||
MaxLines int
|
||
// MaxBytes 本次导入允许的最大字节数，<=0 时取 defaultMaxSourceBytes。
|
||
MaxBytes int
|
||
}
|
||
|
||
// 导入上限。批量任务的来源完全由管理端指定
|
||
const (
|
||
defaultMaxSourceBytes = 32 * 1024 * 1024 // 单份来源最大字节数
|
||
defaultMaxSourceLines = 500000 // 单份来源最大行数
|
||
fetchTimeout = 2 * time.Minute // 远端来源拉取整体超时
|
||
)
|
||
|
||
// 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 通用批量处理函数。
|
||
//
|
||
// 返回 error 而不是像以前那样只打日志静默返回：调用方（手工执行接口 / 定时任务）
|
||
// 需要据此给用户明确结果并写安全审计。以前无论来源读不出、执行方式非法还是
|
||
// 一条都没导进去，接口都回"手工执行任务成功"，运维完全看不出任务其实没跑。
|
||
func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config BatchProcessorConfig) error {
|
||
innerLogName := config.LogPrefix
|
||
|
||
// 枚举在这里再兜一次：保存侧已经拦了，但存量库里可能留着旧的非法值。
|
||
// 执行方式尤其要紧——各 processor 对未知值一律 return false，
|
||
// 表现为"跑完了、什么也没发生"，从外面完全看不出来。
|
||
if !enums.IsValidBatchExecuteMethod(task.BatchExecuteMethod) {
|
||
err := fmt.Errorf("任务执行方式非法(%s)，仅支持 %s/%s",
|
||
task.BatchExecuteMethod, enums.BATCHTASK_EXECUTEMETHODAPPEND, enums.BATCHTASK_EXECUTEMETHODOVERWRITE)
|
||
zlog.Error(innerLogName, err.Error())
|
||
return err
|
||
}
|
||
// 来源类型非法时会落进远端分支（虽然那里也会因为不是合法 URL 而拒），
|
||
// 但报错说"远端来源被拒绝"会把人往错误方向带，不如在这里说清楚。
|
||
if !enums.IsValidBatchSourceType(task.BatchSourceType) {
|
||
err := fmt.Errorf("来源类型非法(%s)，仅支持 %s/%s", task.BatchSourceType,
|
||
enums.BATCHTASK_SOURCETYPE_LOCAL, enums.BATCHTASK_SOURCETYPE_REMOTE)
|
||
zlog.Error(innerLogName, err.Error())
|
||
return err
|
||
}
|
||
|
||
// 一次读取、校验并落到内存：既做上限判定，也顺带消灭了"计数一遍 + 处理一遍"
|
||
// 造成的两次拉取（远端源会被请求两次，且两次内容可能不一致）。
|
||
content, err := loadSource(task, config)
|
||
if err != nil {
|
||
zlog.Error(innerLogName, err.Error())
|
||
return err
|
||
}
|
||
|
||
// 判断数据库是否已经关闭
|
||
if global.GWAF_LOCAL_DB == nil {
|
||
err = errors.New("数据库已经关闭批量处理终止")
|
||
zlog.Error(innerLogName, err.Error())
|
||
return err
|
||
}
|
||
|
||
// 获取对应类型的提取器
|
||
extractor := GetExtractor(task.BatchType)
|
||
|
||
// 先整份扫一遍算行数。扫描出错就直接放弃：这一遍还没往库里写任何东西，
|
||
// 拒绝是干净的；硬着头皮往下走的话，正式那一遍会在同一个位置出错，
|
||
// 而此时前半截已经落库了——就成了不该出现的半截导入。
|
||
totalLines, validLines, err := countLinesWithExtractor(content, extractor)
|
||
if err != nil {
|
||
err = fmt.Errorf("来源扫描失败，已拒绝整批导入: %v", err)
|
||
zlog.Error(innerLogName, err.Error())
|
||
return err
|
||
}
|
||
|
||
// 创建进度跟踪器
|
||
progress := &BatchProgress{
|
||
TotalItems: int32(validLines),
|
||
}
|
||
|
||
zlog.Info(innerLogName, fmt.Sprintf("开始批量处理，总行数: %d，有效行数: %d", totalLines, validLines))
|
||
|
||
// 收集有效的项目
|
||
validItems := make([]string, 0, config.BatchSize)
|
||
scanner := bufio.NewScanner(bytes.NewReader(content))
|
||
hasAffectInfo := false
|
||
|
||
// 按批次处理
|
||
batchCount := 0
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
line = strings.TrimSpace(line)
|
||
if line == "" {
|
||
continue // 跳过空行
|
||
}
|
||
|
||
// 使用特定类型的提取器
|
||
item := extractor.ExtractItem(line)
|
||
if !extractor.ValidateItem(item) {
|
||
continue
|
||
}
|
||
|
||
validItems = append(validItems, item)
|
||
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))
|
||
}
|
||
|
||
// 先判定源是否被完整读完，收尾动作（如全量同步的删除）依赖这个信号
|
||
scanErr := scanner.Err()
|
||
if scanErr != nil {
|
||
zlog.Error(innerLogName, fmt.Sprintf("扫描文件时发生错误: %s", scanErr.Error()))
|
||
}
|
||
|
||
// 收尾回调（可选实现）
|
||
if finalizer, ok := processor.(BatchFinalizer); ok {
|
||
if finalizer.Finalize(task, progress, scanErr == nil) {
|
||
hasAffectInfo = true
|
||
}
|
||
}
|
||
|
||
if hasAffectInfo {
|
||
// 通知引擎进行实时生效
|
||
processor.NotifyEngine(task)
|
||
}
|
||
|
||
// 输出最终处理结果
|
||
zlog.Info(innerLogName, fmt.Sprintf("批量处理完成，总计处理: %d，插入: %d，更新: %d",
|
||
progress.ProcessedItems, progress.InsertedItems, progress.UpdatedItems))
|
||
return scanErr
|
||
}
|
||
|
||
// countLinesWithExtractor 使用提取器计算总行数和有效行数
|
||
func countLinesWithExtractor(content []byte, extractor ItemExtractor) (int, int, error) {
|
||
scanner := bufio.NewScanner(bytes.NewReader(content))
|
||
totalLines := 0
|
||
validLines := 0
|
||
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
totalLines++
|
||
|
||
line = strings.TrimSpace(line)
|
||
if line == "" {
|
||
continue // 跳过空行
|
||
}
|
||
|
||
item := extractor.ExtractItem(line)
|
||
if extractor.ValidateItem(item) {
|
||
validLines++
|
||
}
|
||
}
|
||
|
||
return totalLines, validLines, scanner.Err()
|
||
}
|
||
|
||
// loadSource 打开、校验并读取数据源，返回全部内容。字节数/行数任一超限即拒绝整批。
|
||
func loadSource(task model.BatchTask, config BatchProcessorConfig) ([]byte, error) {
|
||
maxLines := config.MaxLines
|
||
if maxLines <= 0 {
|
||
maxLines = defaultMaxSourceLines
|
||
}
|
||
maxBytes := config.MaxBytes
|
||
if maxBytes <= 0 {
|
||
maxBytes = defaultMaxSourceBytes
|
||
}
|
||
|
||
contentReader, expectedLen, err := openSource(task)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer contentReader.Close()
|
||
|
||
// 多读 1 字节用于判定"是否超限"，而不是静默截断成半份数据
|
||
data, err := io.ReadAll(io.LimitReader(contentReader, int64(maxBytes)+1))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("读取来源失败: %v", err)
|
||
}
|
||
if len(data) > maxBytes {
|
||
return nil, fmt.Errorf("来源超过大小上限 %d 字节，已拒绝整批导入", maxBytes)
|
||
}
|
||
// 声明了长度却没读满 = 中途被掐断。这一步很关键：连接关闭式分帧的 HTTP 响应
|
||
// 被截断时 io.ReadAll 是不报错的，半份数据会被当成"完整的源"，
|
||
// 覆盖模式据此把组里"本次没出现"的条目全删掉——一次网络抖动就能清空 IP 组。
|
||
if expectedLen >= 0 && int64(len(data)) != expectedLen {
|
||
return nil, fmt.Errorf("来源读取不完整(已读 %d 字节，声明 %d 字节)，已拒绝整批导入", len(data), expectedLen)
|
||
}
|
||
if lines := countRawLines(data); lines > maxLines {
|
||
return nil, fmt.Errorf("来源行数 %d 超过上限 %d，已拒绝整批导入", lines, maxLines)
|
||
}
|
||
// bufio.Scanner 单行上限 64KB，超了会在扫描中途报错。放到这里判：
|
||
// 此时一条都还没落库，可以干净地拒绝整批；等扫描时才发现的话，
|
||
// 长行之前的数据已经写进去了，就成了半截导入。
|
||
if n := maxLineLen(data); n > bufio.MaxScanTokenSize {
|
||
return nil, fmt.Errorf("来源存在超长行(%d 字节，上限 %d)，已拒绝整批导入", n, bufio.MaxScanTokenSize)
|
||
}
|
||
return data, nil
|
||
}
|
||
|
||
// maxLineLen 返回最长一行的字节数（不含换行符）
|
||
func maxLineLen(data []byte) int {
|
||
max, cur := 0, 0
|
||
for _, b := range data {
|
||
if b == '\n' {
|
||
if cur > max {
|
||
max = cur
|
||
}
|
||
cur = 0
|
||
continue
|
||
}
|
||
cur++
|
||
}
|
||
if cur > max {
|
||
max = cur
|
||
}
|
||
return max
|
||
}
|
||
|
||
// countRawLines 统计原始行数（末行无换行也算一行）
|
||
func countRawLines(data []byte) int {
|
||
if len(data) == 0 {
|
||
return 0
|
||
}
|
||
n := bytes.Count(data, []byte{'\n'})
|
||
if data[len(data)-1] != '\n' {
|
||
n++
|
||
}
|
||
return n
|
||
}
|
||
|
||
// openSource 打开本地或远程数据源。
|
||
//
|
||
// 这两条路径的入参 BatchSource 完全由管理端提交，必须先过安全边界：
|
||
// - local：只允许读内置 data/import 目录与 config.yml 带外声明的目录（见 utils.ValidateBatchLocalPath），
|
||
// 否则任意绝对路径都能被读出来并逐行落库、再经列表接口读回。
|
||
// - 远端：必须是允许的对外地址（默认仅公网）
|
||
//
|
||
// 破坏性变更：升级前把来源放在任意路径、或指向内网镜像的任务会开始报错，
|
||
// 需要把文件挪进允许目录 / 在 config.yml 的 security.outbound_allowed_hosts 声明该主机。
|
||
// 第二个返回值是"声明的字节数"，未知时为 -1，供调用方判定读取是否完整。
|
||
func openSource(task model.BatchTask) (io.ReadCloser, int64, error) {
|
||
if enums.IsBatchLocalSource(task.BatchSourceType) {
|
||
realPath, err := utils.ValidateBatchLocalPath(task.BatchSource)
|
||
if err != nil {
|
||
return nil, -1, fmt.Errorf("本地来源被拒绝: %v", err)
|
||
}
|
||
f, err := os.Open(realPath)
|
||
if err != nil {
|
||
return nil, -1, fmt.Errorf("failed to open local file: %v", err)
|
||
}
|
||
size := int64(-1)
|
||
if fi, statErr := f.Stat(); statErr == nil {
|
||
size = fi.Size()
|
||
}
|
||
return f, size, nil
|
||
}
|
||
// remote
|
||
rawURL := strings.TrimSpace(task.BatchSource)
|
||
if ok, reason := utils.IsAllowedOutboundURL(rawURL); !ok {
|
||
return nil, -1, fmt.Errorf("远端来源被拒绝: %s", reason)
|
||
}
|
||
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
|
||
if err != nil {
|
||
return nil, -1, fmt.Errorf("failed to build request: %v", err)
|
||
}
|
||
req.Header.Set("User-Agent", "SamWaf-BatchTask")
|
||
resp, err := utils.SafeOutboundHTTPClient(fetchTimeout).Do(req)
|
||
if err != nil {
|
||
return nil, -1, fmt.Errorf("failed to fetch remote data: %v", err)
|
||
}
|
||
if resp.StatusCode != http.StatusOK {
|
||
resp.Body.Close()
|
||
return nil, -1, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
|
||
}
|
||
return resp.Body, resp.ContentLength, 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 // 如果没有匹配到，返回原始行
|
||
}
|