From 0ce0c668b95add1a25cb0026497edc006e5c4f99 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Fri, 11 Mar 2022 10:50:15 +0300
Subject: [PATCH] Fix broken whitelists

---
 internal/cli/run_proxy.go | 8 +++++---
 ipblocklist/firehol.go    | 2 +-
 ipblocklist/noop.go       | 5 ++++-
 ipblocklist/noop_test.go  | 7 +++++++
 mtglib/init.go            | 6 ++++++
 mtglib/proxy.go           | 6 ++++++
 6 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go
index a53e1d7..eec0fae 100644
--- a/internal/cli/run_proxy.go
+++ b/internal/cli/run_proxy.go
@@ -110,6 +110,8 @@ func makeIPBlocklist(conf config.ListConfig, logger mtglib.Logger, ntw mtglib.Ne
 		return nil, fmt.Errorf("incorrect parameters for firehol: %w", err)
 	}
 
+	go firehol.Run(conf.UpdateEach.Get(ipblocklist.DefaultFireholUpdateEach))
+
 	return firehol, nil
 }
 
@@ -162,7 +164,7 @@ func runProxy(conf *config.Config, version string) error { // nolint: funlen
 		return fmt.Errorf("cannot build network: %w", err)
 	}
 
-	blocklist, err := makeIPBlocklist(conf.Defense.Blocklist, logger, ntw)
+	blocklist, err := makeIPBlocklist(conf.Defense.Blocklist, logger.Named("blocklist"), ntw)
 	if err != nil {
 		return fmt.Errorf("cannot build ip blocklist: %w", err)
 	}
@@ -170,9 +172,9 @@ func runProxy(conf *config.Config, version string) error { // nolint: funlen
 	var whitelist mtglib.IPBlocklist
 
 	if conf.Defense.Allowlist.Enabled.Get(false) {
-		whlist, err := makeIPBlocklist(conf.Defense.Allowlist, logger, ntw)
+		whlist, err := makeIPBlocklist(conf.Defense.Allowlist, logger.Named("allowlist"), ntw)
 		if err != nil {
-			return fmt.Errorf("cannot build ip blocklist: %w", err)
+			return fmt.Errorf("cannot build ip allowlist: %w", err)
 		}
 
 		whitelist = whlist
diff --git a/ipblocklist/firehol.go b/ipblocklist/firehol.go
index 41726a7..05d34cb 100644
--- a/ipblocklist/firehol.go
+++ b/ipblocklist/firehol.go
@@ -163,7 +163,7 @@ func (f *Firehol) update() {
 	f.treeV4 = v4tree
 	f.treeV6 = v6tree
 
-	f.logger.Info("blocklist was updated")
+	f.logger.Info("ip list was updated")
 }
 
 func (f *Firehol) updateFromFile(mutex sync.Locker,
diff --git a/ipblocklist/noop.go b/ipblocklist/noop.go
index 2df31cb..40da20f 100644
--- a/ipblocklist/noop.go
+++ b/ipblocklist/noop.go
@@ -2,13 +2,16 @@ package ipblocklist
 
 import (
 	"net"
+	"time"
 
 	"github.com/9seconds/mtg/v2/mtglib"
 )
 
 type noop struct{}
 
-func (n noop) Contains(ip net.IP) bool { return false }
+func (n noop) Contains(ip net.IP) bool      { return false }
+func (n noop) Run(updateEach time.Duration) {}
+func (n noop) Shutdown()                    {}
 
 // NewNoop returns a dummy ipblocklist which allows all incoming
 // connections.
diff --git a/ipblocklist/noop_test.go b/ipblocklist/noop_test.go
index f5bb346..31c6327 100644
--- a/ipblocklist/noop_test.go
+++ b/ipblocklist/noop_test.go
@@ -17,6 +17,13 @@ func (suite *NoopTestSuite) TestOp() {
 	suite.False(ipblocklist.NewNoop().Contains(net.ParseIP("10.0.0.10")))
 }
 
+func (suite *NoopTestSuite) TestRun() {
+	blocklist := ipblocklist.NewNoop()
+
+	blocklist.Run(0)
+	blocklist.Shutdown()
+}
+
 func TestNoop(t *testing.T) {
 	t.Parallel()
 	suite.Run(t, &NoopTestSuite{})
diff --git a/mtglib/init.go b/mtglib/init.go
index 2f34ff7..a245935 100644
--- a/mtglib/init.go
+++ b/mtglib/init.go
@@ -176,6 +176,12 @@ type IPBlocklist interface {
 	// Contains checks if given IP address belongs to this blocklist If.
 	// it is, a connection is terminated                               .
 	Contains(net.IP) bool
+
+	// Run starts a background update procedure for a blocklist
+	Run(time.Duration)
+
+	// Shutdown stops a blocklist. It is assumed that none will access it after.
+	Shutdown()
 }
 
 // Event is a data structure which is populated during mtg request
diff --git a/mtglib/proxy.go b/mtglib/proxy.go
index 294d55d..dea9aa1 100644
--- a/mtglib/proxy.go
+++ b/mtglib/proxy.go
@@ -144,6 +144,12 @@ func (p *Proxy) Shutdown() {
 	p.ctxCancel()
 	p.streamWaitGroup.Wait()
 	p.workerPool.Release()
+
+	if p.whitelist != nil {
+		p.whitelist.Shutdown()
+	}
+
+	p.blocklist.Shutdown()
 }
 
 func (p *Proxy) doFakeTLSHandshake(ctx *streamContext) bool {