feat: HTTP FLOOD Speed Limit

This commit is contained in:
Roi Feng
2025-04-09 21:47:47 -04:00
parent b49c3573bb
commit 4893e95bb7
3 changed files with 34 additions and 2 deletions

View File

@ -165,3 +165,16 @@ func (tc *Counter) GC() {
bucket.mu.Unlock() bucket.mu.Unlock()
} }
} }
func StartCounterGC(counter *Counter, interval time.Duration, stopCh <-chan struct{}) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
counter.GC()
case <-stopCh:
return
}
}
}

View File

@ -39,3 +39,13 @@ func ParseRate(s string) (int64, int64, error) {
} }
return int64(limit), int64(seconds), nil return int64(limit), int64(seconds), nil
} }
func FindMaxRateTime(rateList map[int64]int64) int64 {
maxTimeWindow := int64(0)
for window := range rateList {
if window > maxTimeWindow {
maxTimeWindow = window
}
}
return maxTimeWindow
}

13
main.go
View File

@ -5,10 +5,13 @@ import (
"log" "log"
"os" "os"
"os/signal" "os/signal"
"runtime"
"server_torii/internal/config" "server_torii/internal/config"
"server_torii/internal/dataType" "server_torii/internal/dataType"
"server_torii/internal/server" "server_torii/internal/server"
"server_torii/internal/utils"
"syscall" "syscall"
"time"
) )
func main() { func main() {
@ -49,10 +52,15 @@ func main() {
//allocate shared memory //allocate shared memory
sharedMem := &dataType.SharedMemory{ sharedMem := &dataType.SharedMemory{
HTTPFloodSpeedLimitCounter: dataType.NewCounter(64, 60), HTTPFloodSpeedLimitCounter: dataType.NewCounter(max(runtime.NumCPU()*8, 16), utils.FindMaxRateTime(ruleSet.HTTPFloodRule.HTTPFloodSpeedLimit)),
HTTPFloodSameURILimitCounter: dataType.NewCounter(64, 60), HTTPFloodSameURILimitCounter: dataType.NewCounter(max(runtime.NumCPU()*8, 16), utils.FindMaxRateTime(ruleSet.HTTPFloodRule.HTTPFloodSameURILimit)),
} }
//GC
gcStopCh := make(chan struct{})
go dataType.StartCounterGC(sharedMem.HTTPFloodSpeedLimitCounter, time.Minute, gcStopCh)
go dataType.StartCounterGC(sharedMem.HTTPFloodSameURILimitCounter, time.Minute, gcStopCh)
// Start server // Start server
stop := make(chan os.Signal, 1) stop := make(chan os.Signal, 1)
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
@ -65,6 +73,7 @@ func main() {
select { select {
case <-stop: case <-stop:
log.Println("Stopping server...") log.Println("Stopping server...")
close(gcStopCh)
case err := <-serverErr: case err := <-serverErr:
if err != nil { if err != nil {
log.Fatalf("Failed to start server: %v", err) log.Fatalf("Failed to start server: %v", err)