Merge pull request #238 from samwafgo/feat_cc_ipmode

feat:add cc ip model

api/waf_anticc.go

@@ -31,7 +31,7 @@ func (w *WafAntiCCApi) AddApi(c *gin.Context) {
 				response.OkWithMessage("添加成功", c)
 			} else {
 
-				response.FailWithMessage("添加失败", c)
+				response.FailWithMessage("添加失败 "+err.Error(), c)
 			}
 			return
 		} else {

model/anticc.go

@@ -10,12 +10,12 @@ import "SamWaf/model/baseorm"
 
 type AntiCC struct {
 	baseorm.BaseOrm
-	Id            string `json:"id" gorm:"column:id;primary_key"`
 	HostCode      string `json:"host_code" gorm:"column:host_code"`
 	Rate          int    `json:"rate" gorm:"column:rate"`
 	Limit         int    `json:"limit" gorm:"column:limit"`
 	LockIPMinutes int    `json:"lock_ip_minutes" gorm:"column:lock_ip_minutes"`
 	LimitMode     string `json:"limit_mode" gorm:"column:limit_mode"` // "rate" 或 "window"
+	IPMode        string `json:"ip_mode" gorm:"column:ip_mode"`       // "nic" 网卡模式 或 "proxy" 代理模式
 	Url           string `json:"url"`                                 //保护的url
 	Remarks       string `json:"remarks"`                             //备注
 }

model/request/waf_anticc_req.go

@@ -9,6 +9,7 @@ type WafAntiCCAddReq struct {
 	Limit         int    `json:"limit" form:"limit"`                     //限制
 	LockIPMinutes int    `json:"lock_ip_minutes" form:"lock_ip_minutes"` //封禁分钟
 	LimitMode     string `json:"limit_mode"  form:"limit_mode"`          // "rate" 或 "window"
+	IPMode        string `json:"ip_mode"  form:"ip_mode"`                // "nic" 网卡模式 或 "proxy" 代理模式
 	Remarks       string `json:"remarks" form:"remarks"`                 //备注
 }
 type WafAntiCCSearchReq struct {
@@ -29,6 +30,7 @@ type WafAntiCCEditReq struct {
 	Limit         int    `json:"limit" form:"limit"`                      //限制
 	LockIPMinutes int    `json:"lock_ip_minutes"  form:"lock_ip_minutes"` //封禁分钟
 	LimitMode     string `json:"limit_mode"  form:"limit_mode"`           // "rate" 或 "window"
+	IPMode        string `json:"ip_mode"  form:"ip_mode"`                 // "nic" 网卡模式 或 "proxy" 代理模式
 	Remarks       string `json:"remarks" form:"remarks"`                  //备注
 }
 

service/waf_service/waf_anticc.go

@@ -16,6 +16,12 @@ type WafAntiCCService struct{}
 var WafAntiCCServiceApp = new(WafAntiCCService)
 
 func (receiver *WafAntiCCService) AddApi(req request.WafAntiCCAddReq) error {
+	var existingRecord model.AntiCC
+	result := global.GWAF_LOCAL_DB.Where("host_code = ?", req.HostCode).First(&existingRecord)
+	if result.Error == nil {
+		// 记录已存在，返回错误
+		return errors.New("当前网站已存在CC防护配置")
+	}
 	var bean = &model.AntiCC{
 		BaseOrm: baseorm.BaseOrm{
 			Id:          uuid.NewV4().String(),
@@ -31,6 +37,7 @@ func (receiver *WafAntiCCService) AddApi(req request.WafAntiCCAddReq) error {
 		Url:           req.Url,
 		Remarks:       req.Remarks,
 		LimitMode:     req.LimitMode,
+		IPMode:        req.IPMode,
 	}
 	global.GWAF_LOCAL_DB.Create(bean)
 	return nil
@@ -56,6 +63,7 @@ func (receiver *WafAntiCCService) ModifyApi(req request.WafAntiCCEditReq) error
 		"Remarks":       req.Remarks,
 		"UPDATE_TIME":   customtype.JsonTime(time.Now()),
 		"LimitMode":     req.LimitMode,
+		"IPMode":        req.IPMode,
 	}
 	err := global.GWAF_LOCAL_DB.Model(model.AntiCC{}).Where("id = ?", req.Id).Updates(ipWhiteMap).Error
 

wafdb/patch_sql.go

@@ -144,6 +144,15 @@ func pathCoreSql(db *gorm.DB) {
 	} else {
 		zlog.Info("db", "anti_ccs :limit_mode init successfully")
 	}
+
+	//20250401 CC防护IP提取模式，默认为网卡模式
+	err = db.Exec("UPDATE anti_ccs SET ip_mode='nic' WHERE ip_mode IS NULL ").Error
+	if err != nil {
+		panic("failed to anti_ccs :ip_mode " + err.Error())
+	} else {
+		zlog.Info("db", "anti_ccs :ip_mode init successfully")
+	}
+
 	// 记录结束时间并计算耗时
 	duration := time.Since(startTime)
 	zlog.Info("create core index completely", "duration", duration.String())

wafenginecore/checkcc.go

@@ -22,28 +22,46 @@ func (waf *WafEngine) CheckCC(r *http.Request, weblogbean *innerbean.WebLog, for
 		Title:           "",
 		Content:         "",
 	}
-	// cc 防护 (局部检测 )
+	// cc 防护 (局部检测)
 	if hostTarget.PluginIpRateLimiter != nil {
+		// 根据IP模式选择使用的IP
+		var clientIP string
+		if hostTarget.AntiCCBean.IPMode == "proxy" {
+			clientIP = weblogbean.SRC_IP
+		} else {
+			// 默认使用网卡模式
+			clientIP = weblogbean.NetSrcIp
+		}
 
-		if !hostTarget.PluginIpRateLimiter.Allow(weblogbean.NetSrcIp) {
+		if !hostTarget.PluginIpRateLimiter.Allow(clientIP) {
 			weblogbean.RISK_LEVEL = 1
 			result.IsBlock = true
 			result.Title = "【局部】触发IP频次访问限制"
 			result.Content = "您的访问被阻止超量了"
-			cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.NetSrcIp
+			cacheKey := enums.CACHE_CCVISITBAN_PRE + clientIP
 			//将该IP添加到封禁里
 			global.GCACHE_WAFCACHE.SetWithTTl(cacheKey, 1, time.Duration(hostTarget.AntiCCBean.LockIPMinutes)*time.Minute)
 			return result
 		}
 	}
-	// cc 防护 （全局检测 ）
+
+	// cc 防护 （全局检测）
 	if waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].Host.GUARD_STATUS == 1 && waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter != nil {
-		if !waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter.Allow(weblogbean.NetSrcIp) {
+		// 根据IP模式选择使用的IP
+		var clientIP string
+		if waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].AntiCCBean.IPMode == "proxy" {
+			clientIP = weblogbean.SRC_IP
+		} else {
+			// 默认使用网卡模式
+			clientIP = weblogbean.NetSrcIp
+		}
+
+		if !waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter.Allow(clientIP) {
 			weblogbean.RISK_LEVEL = 1
 			result.IsBlock = true
 			result.Title = "【全局】触发IP频次访问限制"
 			result.Content = "您的访问被阻止超量了"
-			cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.NetSrcIp
+			cacheKey := enums.CACHE_CCVISITBAN_PRE + clientIP
 			//将该IP添加到封禁里
 			global.GCACHE_WAFCACHE.SetWithTTl(cacheKey, 1, time.Duration(waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].AntiCCBean.LockIPMinutes)*time.Minute)
 			return result