included all map access inside locks, add blocked ip list

main.go

@@ -17,7 +17,13 @@ const SameCookieTimeLimitMins = 10
 
 type Ip2LastSeen map[string]time.Time
 
-var cookiesToIps sync.Map // map [ cookie string] Ip2LastSeen
+var cookiesToIps map[string]Ip2LastSeen // map [ cookie string] Ip2LastSeen
+var cookiesLock sync.Mutex
+
+var blocklistedIps map[string]bool
+var blocklistLock sync.Mutex
+
+
 var counter *Counter
 
 func main() {
@@ -36,6 +42,9 @@ func main() {
 	}
 	log.Infof("Starting ddosFilter on %v -> %v...\n", *listenPort, *proxyPort)
 
+	cookiesToIps = map[string]Ip2LastSeen{}
+	blocklistedIps = map[string]bool{}
+
 	counter = NewCounter()
 	go logger()
 	listen(*listenPort, *proxyPort)
@@ -68,6 +77,16 @@ func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort in
 		ip = realIp
 	}
 
+	blocklistLock.Lock()
+	blocked, ok := blocklistedIps[ip]
+	blocklistLock.Unlock()
+	if ok && blocked {
+		log.Debugln("blocked ip, 404ing")
+		res.WriteHeader(http.StatusNotFound)
+		fmt.Fprint(res, "404 - suspected botnet")
+		return
+	}
+
 	//log.Debugf("%v: Request %v %v\n", ip, req.Host, req.URL)
 	cookieObj, err := req.Cookie("i_like_gogits")
 	log.Debugf("ip: %v cookie: %v\n", ip, cookieObj)
@@ -78,41 +97,45 @@ func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort in
 
 	cookie := cookieObj.Value
 
-	ips, ok := cookiesToIps.Load(cookie)
+	cookiesLock.Lock()
+	ips, ok := cookiesToIps[cookie]
+	cookiesLock.Unlock()
+
 	if !ok {
 		ips := Ip2LastSeen{ip: time.Now()}
-		cookiesToIps.Store(cookie, ips)
+		cookiesLock.Lock()
+		cookiesToIps[cookie] = ips
+		cookiesLock.Unlock()
 		pass(res, req, listenPort, proxyPort)
 		return
 	}
 
-	ipsMap := ips.(Ip2LastSeen)
-	if len(ipsMap) > 1 {
-		log.Info("More than 1 IP for cookie: %v %v\n", cookie, ipsMap)
-	}
 	var mostRecent string = ""
-
-	for ip, lastSeen := range ipsMap {
-		if mostRecent == "" || lastSeen.After(ipsMap[mostRecent]) {
+	cookiesLock.Lock()
+	for ip, lastSeen := range ips {
+		if mostRecent == "" || lastSeen.After(ips[mostRecent]) {
 			mostRecent = ip
 		}
 	}
-
-	ipsMap[ip] = time.Now()
-	cookiesToIps.Store(cookie, ipsMap)
+	ips[ip] = time.Now()
+	cookiesToIps[cookie] =  ips
+	cookiesLock.Unlock()
 
 	if mostRecent == "" || mostRecent == ip {
 		pass(res, req, listenPort, proxyPort)
 		return
 	}
 
-	timeDiff := time.Now().Sub(ipsMap[mostRecent])
+	timeDiff := time.Now().Sub(ips[mostRecent])
 	if timeDiff.Minutes() > SameCookieTimeLimitMins {
 		pass(res, req, listenPort, proxyPort)
 		return
 	}
 
 	log.Infof("different IP in the last %v minutes, 404ing\n", SameCookieTimeLimitMins)
+	blocklistLock.Lock()
+	blocklistedIps[ip] = true
+	blocklistLock.Unlock()
 	res.WriteHeader(http.StatusNotFound)
 	fmt.Fprint(res, "404 - suspected botnet")
 }