From 97d29a98b52021a5eec6da16ee56b345c7edb907 Mon Sep 17 00:00:00 2001
From: mphan6 <mphan6@myseneca.ca>
Date: Sat, 8 Dec 2018 23:39:10 -0500
Subject: [PATCH] Whitelist

Added whitelist only functionality through program
---
 api/server.go       | 45 +++++++++++++++++++++++++++++++++++++++++++++
 config.example.json |  3 ++-
 policy/policy.go    |  4 ++++
 storage/redis.go    |  8 ++++++++
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/api/server.go b/api/server.go
index dd021a1e3..97fe25e01 100644
--- a/api/server.go
+++ b/api/server.go
@@ -10,6 +10,9 @@ import (
 	"sync/atomic"
 	"time"
 
+	"fmt"
+	"regexp"
+
 	"github.com/gorilla/mux"
 
 	"github.com/sammy007/open-ethereum-pool/storage"
@@ -45,6 +48,10 @@ type Entry struct {
 	updatedAt int64
 }
 
+type WhiteList struct {
+	IPs []string `json:"ips"`
+}
+
 func NewApiServer(cfg *ApiConfig, backend *storage.RedisClient) *ApiServer {
 	hashrateWindow := util.MustParseDuration(cfg.HashrateWindow)
 	hashrateLargeWindow := util.MustParseDuration(cfg.HashrateLargeWindow)
@@ -108,6 +115,7 @@ func (s *ApiServer) listen() {
 	r.HandleFunc("/api/blocks", s.BlocksIndex)
 	r.HandleFunc("/api/payments", s.PaymentsIndex)
 	r.HandleFunc("/api/accounts/{login:0x[0-9a-fA-F]{40}}", s.AccountIndex)
+	r.HandleFunc("/api/whitelist", s.AddToWhiteListIndex)
 	r.NotFoundHandler = http.HandlerFunc(notFound)
 	err := http.ListenAndServe(s.config.Listen, r)
 	if err != nil {
@@ -115,6 +123,43 @@ func (s *ApiServer) listen() {
 	}
 }
 
+func (s *ApiServer) AddToWhiteListIndex(w http.ResponseWriter, r *http.Request) {
+	if r.Method == "POST" {
+		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
+		w.Header().Set("Access-Control-Allow-Origin", "*")
+		w.Header().Set("Cache-Control", "no-cache")
+
+		reply := make(map[string]interface{})
+		decoder := json.NewDecoder(r.Body)
+		var whiteList WhiteList
+		err := decoder.Decode(&whiteList)
+
+		if err != nil {
+			reply["success"] = false
+			reply["message"] = fmt.Sprintf("Error occured: %v", err)
+            err := json.NewEncoder(w).Encode(reply)
+            if err != nil {
+                log.Println("Error serializing API response: ", err)
+            }
+            w.WriteHeader(http.StatusUnprocessableEntity)
+            return
+		}
+		reg, _ := regexp.Compile("\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}")
+		for _, ip := range whiteList.IPs {
+			if reg.MatchString(ip){
+				err = s.backend.AddToWhiteList(ip)
+				if err != nil {
+					log.Println(err)
+				}
+			}
+		}
+		reply["success"] = true
+		reply["message"] = fmt.Sprintf("Add to IP white list successful.")
+	}else{
+		http.Error(w, "Invalid request method.", http.StatusMethodNotAllowed)
+	}
+}
+
 func notFound(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
 	w.Header().Set("Access-Control-Allow-Origin", "*")
diff --git a/config.example.json b/config.example.json
index 1a264b8cc..74810cd56 100644
--- a/config.example.json
+++ b/config.example.json
@@ -42,7 +42,8 @@
 				"limit": 30,
 				"grace": "5m",
 				"limitJump": 10
-			}
+			},
+			"whiteListOnly": true
 		}
 	},
 
diff --git a/policy/policy.go b/policy/policy.go
index 0e9fcf392..45a5758d3 100644
--- a/policy/policy.go
+++ b/policy/policy.go
@@ -19,6 +19,7 @@ type Config struct {
 	Limits          Limits  `json:"limits"`
 	ResetInterval   string  `json:"resetInterval"`
 	RefreshInterval string  `json:"refreshInterval"`
+	WhiteListOnly   bool    `json:"whiteListOnly"`
 }
 
 type Limits struct {
@@ -186,6 +187,9 @@ func (s *PolicyServer) BanClient(ip string) {
 }
 
 func (s *PolicyServer) IsBanned(ip string) bool {
+	if(s.config.whiteListOnly){
+		return !s.InWhiteList(ip)
+	}
 	x := s.Get(ip)
 	return atomic.LoadInt32(&x.Banned) > 0
 }
diff --git a/storage/redis.go b/storage/redis.go
index 449b58fcc..1d77ad8da 100644
--- a/storage/redis.go
+++ b/storage/redis.go
@@ -957,3 +957,11 @@ func convertPaymentsResults(raw *redis.ZSliceCmd) []map[string]interface{} {
 	}
 	return result
 }
+
+func (r *RedisClient) AddToWhiteList(ip string) error {
+	cmd := r.client.SAdd(r.formatKey("whitelist"), ip)
+	if cmd.Err() != nil {
+		return cmd.Err()
+	}
+	return nil
+}