[i] deferred removal from whitelist

cmd/kf2-antiddos/args.go

@@ -20,6 +20,7 @@ func printHelp() {
 	output.Println("  -j, --jobs N             allow N jobs at once")
 	output.Println("  -o, --output MODE        self|proxy|all|quiet")
 	output.Println("  -t, --deny-time TIME     minimum ip deny TIME (seconds)")
+	output.Println("  -a, --allow-time TIME    ip whitelist time after disconnect (seconds)")
 	output.Println("  -c, --max-connections N  Skip N connections before run deny script")
 	output.Println("  -v, --version            Show version")
 	output.Println("  -h, --help               Show help")
@@ -41,6 +42,9 @@ func parseArgs() config.Config {
 	gnuflag.UintVar(&rawCfg.DenyTime, "t", 0, "")
 	gnuflag.UintVar(&rawCfg.DenyTime, "deny-time", 0, "")
 
+	gnuflag.UintVar(&rawCfg.AllowTime, "a", 0, "")
+	gnuflag.UintVar(&rawCfg.AllowTime, "allow-time", 0, "")
+
 	gnuflag.UintVar(&rawCfg.MaxConn, "c", 0, "")
 	gnuflag.UintVar(&rawCfg.MaxConn, "max-connections", 0, "")
 

cmd/kf2-antiddos/main.go

@@ -95,6 +95,7 @@ func main() {
 			&banChan,
 			&resetChan,
 			cfg.MaxConn,
+			cfg.AllowTime,
 		))
 
 	// Action worker

internal/action/action.go

@@ -60,7 +60,7 @@ func (a *Action) allow(unbanAll bool) {
 	unban := make([]string, 0)
 
 	for ip := range a.ips {
-		if unbanAll || bool(a.ips[ip]) { // aka if readyToUnban
+		if unbanAll || a.ips[ip] { // aka if readyToUnban
 			unban = append(unban, ip)
 		} else {
 			a.ips[ip] = true // mark readyToUnban next time

internal/config/config.go

@@ -22,6 +22,7 @@ type Config struct {
 	Jobs        uint
 	OutputMode  string
 	DenyTime    uint
+	AllowTime   uint
 	MaxConn     uint
 
 	ShowVersion bool
@@ -79,4 +80,7 @@ func (cfg *Config) SetEmptyArgs() {
 	if cfg.DenyTime == 0 {
 		cfg.DenyTime = 20 * 60
 	}
+	if cfg.AllowTime == 0 {
+		cfg.AllowTime = 20 * 60
+	}
 }

internal/history/history.go

@@ -2,9 +2,11 @@ package history
 
 import (
 	"kf2-antiddos/internal/common"
+	"time"
 )
 
 type History struct {
+	ticker    *time.Ticker
 	quit      chan struct{}
 	eventChan *chan common.Event
 	banChan   *chan string
@@ -12,18 +14,19 @@ type History struct {
 	head      byte
 	history   map[byte]common.Event
 	ips       map[string]uint // map[ip]conn_count
-	whitelist map[string]struct{}
+	whitelist map[string]bool
 	banned    map[string]struct{}
 	maxConn   uint
 	workerID  uint
 }
 
-func New(workerID uint, eventChan *chan common.Event, banChan *chan string, resetChan *chan string, maxConn uint) *History {
+func New(workerID uint, eventChan *chan common.Event, banChan *chan string, resetChan *chan string, maxConn uint, allowTime uint) *History {
 	return &History{
+		ticker:    time.NewTicker(time.Duration(allowTime) * time.Second),
 		quit:      make(chan struct{}),
 		ips:       make(map[string]uint, 0),
 		history:   make(map[byte]common.Event, 0),
-		whitelist: make(map[string]struct{}, 0),
+		whitelist: make(map[string]bool, 0),
 		banned:    make(map[string]struct{}, 0),
 		eventChan: eventChan,
 		banChan:   banChan,
@@ -42,7 +45,10 @@ func (h *History) Do() {
 				h.registerEvent(event)
 			case ip := <-*h.resetChan:
 				h.resetIp(ip)
+			case <-h.ticker.C:
+				h.unWhiteList()
 			case <-h.quit:
+				h.ticker.Stop()
 				return
 			}
 		}
@@ -88,16 +94,29 @@ func (h *History) registerConnect(ip string) {
 }
 
 func (h *History) registerNewPlayer(ip string) {
-	h.whitelist[ip] = struct{}{}
+	h.whitelist[ip] = false
 }
 
 func (h *History) registerEndPlayer(ip string) {
-	delete(h.whitelist, ip)
+	h.whitelist[ip] = true
-	delete(h.ips, ip)
-	delete(h.banned, ip)
 }
 
 func (h *History) resetIp(ip string) {
 	delete(h.ips, ip)
 	delete(h.banned, ip)
 }
+
+func (h *History) unWhiteList() {
+	toRemove := make([]string, 0)
+	for ip := range h.whitelist {
+		if h.whitelist[ip] {
+			toRemove = append(toRemove, ip)
+		}
+	}
+
+	for _, ip := range toRemove {
+		delete(h.whitelist, ip)
+		delete(h.ips, ip)
+		delete(h.banned, ip)
+	}
+}