mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-02-04 12:40:07 +00:00
refactor: ip data type
This commit is contained in:
@@ -2,18 +2,14 @@ package waf
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"errors"
|
||||
"log"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"github.com/naiba/nezha/model"
|
||||
"github.com/naiba/nezha/service/singleton"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
//go:embed waf.html
|
||||
@@ -55,33 +51,13 @@ func Waf(c *gin.Context) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
var w model.WAF
|
||||
if err := singleton.DB.First(&w, "ip = ?", realipAddr).Error; err != nil {
|
||||
if err != gorm.ErrRecordNotFound {
|
||||
ShowBlockPage(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
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))
|
||||
ShowBlockPage(c, errors.New("you are blocked by nezha WAF"))
|
||||
if err := model.CheckIP(singleton.DB, realipAddr); err != nil {
|
||||
ShowBlockPage(c, err)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
|
||||
func pow(x, y uint64) uint64 {
|
||||
base := big.NewInt(0).SetUint64(x)
|
||||
exp := big.NewInt(0).SetUint64(y)
|
||||
result := big.NewInt(1)
|
||||
result.Exp(base, exp, nil)
|
||||
if !result.IsUint64() {
|
||||
return ^uint64(0) // return max uint64 value on overflow
|
||||
}
|
||||
return result.Uint64()
|
||||
}
|
||||
|
||||
func ShowBlockPage(c *gin.Context, err error) {
|
||||
c.Writer.WriteHeader(http.StatusForbidden)
|
||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
package waf
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPow(t *testing.T) {
|
||||
tests := []struct {
|
||||
x,
|
||||
y,
|
||||
expect uint64
|
||||
}{
|
||||
{2, 64, math.MaxUint64}, // 2 的 64 次方,超过 uint64 最大值
|
||||
{uint64(1 << 63), 2, math.MaxUint64}, // 大数平方,可能溢出
|
||||
{uint64(^uint64(0)), 2, math.MaxUint64}, // uint64 最大值的平方,溢出
|
||||
{2, 3, 8},
|
||||
{5, 0, 1},
|
||||
{3, 1, 3},
|
||||
{0, 5, 0},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
result := pow(tt.x, tt.y)
|
||||
if result != tt.expect {
|
||||
t.Errorf("pow(%d, %d) = %d; expect %d", tt.x, tt.y, result, tt.expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user