Merge pull request #118 from creekorful/106-improve-blacklister

Implement new blacklister

deployments/docker/docker-compose.yml

@@ -85,6 +85,7 @@ services:
       --default-value forbidden-hostnames="[{\"hostname\": \"facebookcorewwwi.onion\"}, {\"hostname\": \"nytimes3xbfgragh.onion\"}]"
       --default-value allowed-mime-types="[{\"content-type\":\"text/\",\"extensions\":[\"html\",\"php\",\"aspx\", \"htm\"]}]"
       --default-value refresh-delay="{\"delay\": -1}"
+      --default-value blacklist-threshold="{\"threshold\": 10}"
     restart: always
     depends_on:
       - rabbitmq
@@ -97,10 +98,12 @@ services:
       --log-level debug
       --hub-uri amqp://guest:guest@rabbitmq:5672
       --config-api-uri http://configapi:8080
+      --redis-uri redis:6379
     restart: always
     depends_on:
       - rabbitmq
       - configapi
+      - redis
 
 volumes:
   esdata:

internal/blacklister/blacklister.go

@@ -1,6 +1,8 @@
 package blacklister
 
 import (
+	"fmt"
+	"github.com/creekorful/trandoshan/internal/cache"
 	configapi "github.com/creekorful/trandoshan/internal/configapi/client"
 	"github.com/creekorful/trandoshan/internal/event"
 	"github.com/creekorful/trandoshan/internal/process"
@@ -12,7 +14,8 @@ import (
 
 // State represent the application state
 type State struct {
-	configClient configapi.Client
+	configClient  configapi.Client
+	hostnameCache cache.Cache
 }
 
 // Name return the process name
@@ -22,7 +25,7 @@ func (state *State) Name() string {
 
 // CommonFlags return process common flags
 func (state *State) CommonFlags() []string {
-	return []string{process.HubURIFlag, process.ConfigAPIURIFlag}
+	return []string{process.HubURIFlag, process.ConfigAPIURIFlag, process.RedisURIFlag}
 }
 
 // CustomFlags return process custom flags
@@ -32,7 +35,13 @@ func (state *State) CustomFlags() []cli.Flag {
 
 // Initialize the process
 func (state *State) Initialize(provider process.Provider) error {
-	client, err := provider.ConfigClient([]string{configapi.ForbiddenHostnamesKey})
+	hostnameCache, err := provider.Cache()
+	if err != nil {
+		return err
+	}
+	state.hostnameCache = hostnameCache
+
+	client, err := provider.ConfigClient([]string{configapi.ForbiddenHostnamesKey, configapi.BlackListThresholdKey})
 	if err != nil {
 		return err
 	}
@@ -64,29 +73,52 @@ func (state *State) handleTimeoutURLEvent(subscriber event.Subscriber, msg event
 		return err
 	}
 
-	forbiddenHostnames, err := state.configClient.GetForbiddenHostnames()
+	threshold, err := state.configClient.GetBlackListThreshold()
 	if err != nil {
 		return err
 	}
 
-	// prevent duplicates
-	for _, hostname := range forbiddenHostnames {
-		if hostname.Hostname == u.Hostname() {
+	cacheKey := fmt.Sprintf("hostnames:%s", u.Hostname())
+	count, err := state.hostnameCache.GetInt64(cacheKey)
+	if err != nil && err != cache.ErrNIL {
+		return err
+	}
+	count++
+
+	if count >= threshold.Threshold {
+		log.Info().
+			Str("hostname", u.Hostname()).
+			Int64("count", count).
+			Msg("Blacklisting hostname")
+
+		forbiddenHostnames, err := state.configClient.GetForbiddenHostnames()
+		if err != nil {
+			return err
+		}
+
+		// prevent duplicates
+		found := false
+		for _, hostname := range forbiddenHostnames {
+			if hostname.Hostname == u.Hostname() {
+				found = true
+				break
+			}
+		}
+
+		if found {
 			log.Trace().Str("hostname", u.Hostname()).Msg("skipping duplicate hostname")
-			return nil
+		} else {
+			forbiddenHostnames = append(forbiddenHostnames, configapi.ForbiddenHostname{Hostname: u.Hostname()})
+			if err := state.configClient.Set(configapi.ForbiddenHostnamesKey, forbiddenHostnames); err != nil {
+				return err
+			}
 		}
 	}
 
-	forbiddenHostnames = append(forbiddenHostnames, configapi.ForbiddenHostname{Hostname: u.Hostname()})
-
-	if err := state.configClient.Set(configapi.ForbiddenHostnamesKey, forbiddenHostnames); err != nil {
+	// Update count
+	if err := state.hostnameCache.SetInt64(cacheKey, count, cache.NoTTL); err != nil {
 		return err
 	}
 
-	log.Debug().
-		Str("url", evt.URL).
-		Str("hostname", u.Hostname()).
-		Msg("Blacklisted new hostname")
-
 	return nil
 }

internal/blacklister/blacklister_test.go

@@ -1,6 +1,8 @@
 package blacklister
 
 import (
+	"github.com/creekorful/trandoshan/internal/cache"
+	"github.com/creekorful/trandoshan/internal/cache_mock"
 	configapi "github.com/creekorful/trandoshan/internal/configapi/client"
 	"github.com/creekorful/trandoshan/internal/configapi/client_mock"
 	"github.com/creekorful/trandoshan/internal/event"
@@ -9,12 +11,39 @@ import (
 	"testing"
 )
 
+func TestHandleTimeoutURLEventNoDispatch(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	subscriberMock := event_mock.NewMockSubscriber(mockCtrl)
+	configClientMock := client_mock.NewMockClient(mockCtrl)
+	hostnameCacheMock := cache_mock.NewMockCache(mockCtrl)
+
+	msg := event.RawMessage{}
+	subscriberMock.EXPECT().
+		Read(&msg, &event.TimeoutURLEvent{}).
+		SetArg(1, event.TimeoutURLEvent{
+			URL: "https://down-example.onion",
+		}).Return(nil)
+
+	configClientMock.EXPECT().GetBlackListThreshold().Return(configapi.BlackListThreshold{Threshold: 10}, nil)
+
+	hostnameCacheMock.EXPECT().GetInt64("hostnames:down-example.onion").Return(int64(0), cache.ErrNIL)
+	hostnameCacheMock.EXPECT().SetInt64("hostnames:down-example.onion", int64(1), cache.NoTTL).Return(nil)
+
+	s := State{configClient: configClientMock, hostnameCache: hostnameCacheMock}
+	if err := s.handleTimeoutURLEvent(subscriberMock, msg); err != nil {
+		t.Fail()
+	}
+}
+
 func TestHandleTimeoutURLEvent(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()
 
 	subscriberMock := event_mock.NewMockSubscriber(mockCtrl)
 	configClientMock := client_mock.NewMockClient(mockCtrl)
+	hostnameCacheMock := cache_mock.NewMockCache(mockCtrl)
 
 	msg := event.RawMessage{}
 	subscriberMock.EXPECT().
@@ -23,6 +52,10 @@ func TestHandleTimeoutURLEvent(t *testing.T) {
 			URL: "https://down-example.onion",
 		}).Return(nil)
 
+	configClientMock.EXPECT().GetBlackListThreshold().Return(configapi.BlackListThreshold{Threshold: 10}, nil)
+
+	hostnameCacheMock.EXPECT().GetInt64("hostnames:down-example.onion").Return(int64(9), nil)
+
 	configClientMock.EXPECT().
 		GetForbiddenHostnames().
 		Return([]configapi.ForbiddenHostname{{Hostname: "facebookcorewwwi.onion"}}, nil)
@@ -33,7 +66,9 @@ func TestHandleTimeoutURLEvent(t *testing.T) {
 		}).
 		Return(nil)
 
-	s := State{configClient: configClientMock}
+	hostnameCacheMock.EXPECT().SetInt64("hostnames:down-example.onion", int64(10), cache.NoTTL).Return(nil)
+
+	s := State{configClient: configClientMock, hostnameCache: hostnameCacheMock}
 	if err := s.handleTimeoutURLEvent(subscriberMock, msg); err != nil {
 		t.Fail()
 	}
@@ -45,6 +80,7 @@ func TestHandleTimeoutURLEventNoDuplicates(t *testing.T) {
 
 	subscriberMock := event_mock.NewMockSubscriber(mockCtrl)
 	configClientMock := client_mock.NewMockClient(mockCtrl)
+	hostnameCacheMock := cache_mock.NewMockCache(mockCtrl)
 
 	msg := event.RawMessage{}
 	subscriberMock.EXPECT().
@@ -53,11 +89,19 @@ func TestHandleTimeoutURLEventNoDuplicates(t *testing.T) {
 			URL: "https://facebookcorewwwi.onion",
 		}).Return(nil)
 
+	configClientMock.EXPECT().GetBlackListThreshold().Return(configapi.BlackListThreshold{Threshold: 3}, nil)
+
+	hostnameCacheMock.EXPECT().GetInt64("hostnames:facebookcorewwwi.onion").Return(int64(2), nil)
+
 	configClientMock.EXPECT().
 		GetForbiddenHostnames().
 		Return([]configapi.ForbiddenHostname{{Hostname: "facebookcorewwwi.onion"}}, nil)
+	// Config not updated since hostname is already 'blacklisted'
+	// this may due because of change in threshold
+
+	hostnameCacheMock.EXPECT().SetInt64("hostnames:facebookcorewwwi.onion", int64(3), cache.NoTTL).Return(nil)
 
-	s := State{configClient: configClientMock}
+	s := State{configClient: configClientMock, hostnameCache: hostnameCacheMock}
 	if err := s.handleTimeoutURLEvent(subscriberMock, msg); err != nil {
 		t.Fail()
 	}

internal/configapi/client/client.go

@@ -21,6 +21,8 @@ const (
 	ForbiddenHostnamesKey = "forbidden-hostnames"
 	// RefreshDelayKey is the key to access the refresh delay config
 	RefreshDelayKey = "refresh-delay"
+	// BlackListThresholdKey is the key to access the blacklist threshold config
+	BlackListThresholdKey = "blacklist-threshold"
 )
 
 // MimeType is the mime type as represented in the config
@@ -41,11 +43,17 @@ type RefreshDelay struct {
 	Delay time.Duration `json:"delay"`
 }
 
+// BlackListThreshold is the threshold to reach before blacklisting domain
+type BlackListThreshold struct {
+	Threshold int64 `json:"threshold"`
+}
+
 // Client is a nice client interface for the ConfigAPI
 type Client interface {
 	GetAllowedMimeTypes() ([]MimeType, error)
 	GetForbiddenHostnames() ([]ForbiddenHostname, error)
 	GetRefreshDelay() (RefreshDelay, error)
+	GetBlackListThreshold() (BlackListThreshold, error)
 
 	Set(key string, value interface{}) error
 }
@@ -60,6 +68,7 @@ type client struct {
 	allowedMimeTypes   []MimeType
 	forbiddenHostnames []ForbiddenHostname
 	refreshDelay       RefreshDelay
+	blackListThreshold BlackListThreshold
 }
 
 // NewConfigClient create a new client for the ConfigAPI.
@@ -141,6 +150,22 @@ func (c *client) setRefreshDelay(value RefreshDelay) error {
 	return nil
 }
 
+func (c *client) GetBlackListThreshold() (BlackListThreshold, error) {
+	c.mutexes[BlackListThresholdKey].RLock()
+	defer c.mutexes[BlackListThresholdKey].RUnlock()
+
+	return c.blackListThreshold, nil
+}
+
+func (c *client) setBlackListThreshold(value BlackListThreshold) error {
+	c.mutexes[BlackListThresholdKey].Lock()
+	defer c.mutexes[BlackListThresholdKey].Unlock()
+
+	c.blackListThreshold = value
+
+	return nil
+}
+
 func (c *client) Set(key string, value interface{}) error {
 	b, err := json.Marshal(value)
 	if err != nil {
@@ -207,6 +232,15 @@ func (c *client) setValue(key string, value []byte) error {
 			return err
 		}
 		break
+	case BlackListThresholdKey:
+		var val BlackListThreshold
+		if err := json.Unmarshal(value, &val); err != nil {
+			return err
+		}
+		if err := c.setBlackListThreshold(val); err != nil {
+			return err
+		}
+		break
 	default:
 		return fmt.Errorf("non managed value type: %s", key)
 	}