fix: waf condition

model/waf.go

@@ -2,11 +2,10 @@ package model
 
 import (
 	"errors"
-	"log"
 	"math/big"
-	"net/netip"
 	"time"
 
+	"github.com/naiba/nezha/pkg/utils"
 	"gorm.io/gorm"
 )
 
@@ -27,24 +26,11 @@ func (w *WAF) TableName() string {
 	return "waf"
 }
 
-func ipStringToBinary(ip string) ([]byte, error) {
-	addr, err := netip.ParseAddr(ip)
-	if err != nil {
-		return nil, err
-	}
-	b := addr.As16()
-	return b[:], nil
-}
-
-func binaryToIPString(b []byte) string {
-	var addr16 [16]byte
-	copy(addr16[:], b)
-	addr := netip.AddrFrom16(addr16)
-	return addr.Unmap().String()
-}
-
 func CheckIP(db *gorm.DB, ip string) error {
-	ipBinary, err := ipStringToBinary(ip)
+	if ip == "" {
+		return nil
+	}
+	ipBinary, err := utils.IPStringToBinary(ip)
 	if err != nil {
 		return err
 	}
@@ -55,39 +41,53 @@ func CheckIP(db *gorm.DB, ip string) error {
 		}
 	}
 	now := time.Now().Unix()
-	if w.LastBlockTimestamp+pow(w.Count, 4) > uint64(now) {
-		log.Println(w.Count, w.LastBlockTimestamp+pow(w.Count, 4)-uint64(now))
+	if powAdd(w.Count, 4, w.LastBlockTimestamp) > uint64(now) {
 		return errors.New("you are blocked by nezha WAF")
 	}
 	return nil
 }
 
+func ClearIP(db *gorm.DB, ip string) error {
+	if ip == "" {
+		return nil
+	}
+	ipBinary, err := utils.IPStringToBinary(ip)
+	if err != nil {
+		return err
+	}
+	return db.Delete(&WAF{}, "ip = ?", ipBinary).Error
+}
+
 func BlockIP(db *gorm.DB, ip string, reason uint8) error {
 	if ip == "" {
-		return errors.New("empty ip")
+		return nil
 	}
-	ipBinary, err := ipStringToBinary(ip)
+	ipBinary, err := utils.IPStringToBinary(ip)
 	if err != nil {
 		return err
 	}
 	var w WAF
-	w.LastBlockReason = reason
-	w.LastBlockTimestamp = uint64(time.Now().Unix())
+	w.IP = ipBinary
 	return db.Transaction(func(tx *gorm.DB) error {
-		if err := tx.FirstOrCreate(&w, WAF{IP: ipBinary}).Error; err != nil {
+		if err := tx.Where(&w).Attrs(WAF{
+			LastBlockReason:    reason,
+			LastBlockTimestamp: uint64(time.Now().Unix()),
+		}).FirstOrCreate(&w).Error; err != nil {
 			return err
 		}
-		return tx.Exec("UPDATE waf SET count = count + 1 WHERE ip = ?", ipBinary).Error
+		return tx.Exec("UPDATE waf SET count = count + 1, last_block_reason = ?, last_block_timestamp = ? WHERE ip = ?", reason, uint64(time.Now().Unix()), ipBinary).Error
 	})
 }
 
-func pow(x, y uint64) uint64 {
+func powAdd(x, y, z uint64) uint64 {
 	base := big.NewInt(0).SetUint64(x)
 	exp := big.NewInt(0).SetUint64(y)
 	result := big.NewInt(1)
 	result.Exp(base, exp, nil)
+	result.Add(result, big.NewInt(0).SetUint64(z))
 	if !result.IsUint64() {
 		return ^uint64(0) // return max uint64 value on overflow
 	}
-	return result.Uint64()
+	ret := result.Uint64()
+	return utils.IfOr(ret < z+3, z+3, ret)
 }