cluster code inside on cookie lock to prevent flailing at startup; add saving and loading of blocklist

main.go

@@ -1,11 +1,13 @@
 package main
 
 import (
+	"encoding/json"
 	"flag"
 	"fmt"
 	"git.openprivacy.ca/openprivacy/log"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -14,6 +16,7 @@ import (
 
 const MaxLength = 8192
 const SameCookieTimeLimitMins = 10
+const blocklistFile = "blocklist.json"
 
 type Ip2LastSeen map[string]time.Time
 
@@ -45,11 +48,41 @@ func main() {
 	cookiesToIps = map[string]Ip2LastSeen{}
 	blocklistedIps = map[string]bool{}
 
+	load()
+
 	counter = NewCounter()
 	go logger()
+	go save()
 	listen(*listenPort, *proxyPort)
 }
 
+func load() {
+	_, err := os.Stat(blocklistFile)
+	if os.IsNotExist(err) {
+		return
+	}
+	bytes, err := ioutil.ReadFile(blocklistFile)
+	if err != nil {
+		return
+	}
+	json.Unmarshal(bytes, blocklistedIps)
+	return
+}
+
+func save() {
+	for {
+		time.Sleep(1*time.Minute)
+		blocklistLock.Lock()
+		bytes, err := json.Marshal(blocklistedIps)
+		blocklistLock.Unlock()
+		if err != nil {
+			log.Error("Could not marshal blockedIPlist: %v\n", err)
+			continue
+		}
+		ioutil.WriteFile(blocklistFile, bytes, 0600)
+	}
+}
+
 func logger() {
 	for {
 		time.Sleep(5*time.Second)
@@ -87,7 +120,6 @@ func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort in
 		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)
 	if err != nil {
@@ -97,36 +129,37 @@ func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort in
 
 	cookie := cookieObj.Value
 
-	cookiesLock.Lock()
+	var mostRecent string = ""
-	ips, ok := cookiesToIps[cookie]
+	var ipList = []string{}
-	cookiesLock.Unlock()
+	cookiesLock.Lock() /**** Inside cookies Lock *****/
 
-	if !ok {
+	ipsForCookie, foundIpsForCookie := cookiesToIps[cookie]
+	if !foundIpsForCookie {
 		ips := Ip2LastSeen{ip: time.Now()}
-		cookiesLock.Lock()
 		cookiesToIps[cookie] = ips
-		cookiesLock.Unlock()
+	} else {
+		for ip, lastSeen := range ipsForCookie {
+			ipList = append(ipList, ip)
+			if mostRecent == "" || lastSeen.After(ipsForCookie[mostRecent]) {
+				mostRecent = ip
+			}
+		}
+		ipsForCookie[ip] = time.Now()
+		cookiesToIps[cookie] = ipsForCookie
+	}
+	cookiesLock.Unlock() /***** End Cookies Lock *****/
+
+	if !foundIpsForCookie {
 		pass(res, req, listenPort, proxyPort)
 		return
 	}
 
-	var mostRecent string = ""
-	cookiesLock.Lock()
-	for ip, lastSeen := range ips {
-		if mostRecent == "" || lastSeen.After(ips[mostRecent]) {
-			mostRecent = ip
-		}
-	}
-	ips[ip] = time.Now()
-	cookiesToIps[cookie] =  ips
-	cookiesLock.Unlock()
-
 	if mostRecent == "" || mostRecent == ip {
 		pass(res, req, listenPort, proxyPort)
 		return
 	}
 
-	timeDiff := time.Now().Sub(ips[mostRecent])
+	timeDiff := time.Now().Sub(ipsForCookie[mostRecent])
 	if timeDiff.Minutes() > SameCookieTimeLimitMins {
 		pass(res, req, listenPort, proxyPort)
 		return
@@ -134,7 +167,9 @@ func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort in
 
 	log.Infof("different IP in the last %v minutes, 404ing\n", SameCookieTimeLimitMins)
 	blocklistLock.Lock()
-	blocklistedIps[ip] = true
+	for _, ip := range ipList {
+		blocklistedIps[ip] = true
+	}
 	blocklistLock.Unlock()
 	res.WriteHeader(http.StatusNotFound)
 	fmt.Fprint(res, "404 - suspected botnet")