Merge pull request #231 from 9seconds/whitelists

Whitelist support

example.config.toml

@@ -174,6 +174,27 @@ urls = [
 # How often do we need to update a blocklist set.
 update-each = "24h"
 
+# Allowlist is an opposite to a blocklist. Only those IPs that are coming from
+# subnets defined in these lists are allowed. All others will be rejected.
+#
+# If this feature is disabled, then there won't be any check performed by this
+# validator. It is possible to combine both blocklist and whitelist.
+[defense.allowlist]
+# You can enable/disable this feature.
+enabled = false
+# This is a limiter for concurrency. In order to protect website
+# from overloading, we download files in this number of threads.
+download-concurrency = 2
+# A list of URLs in FireHOL format (https://iplists.firehol.org/)
+# You can provider links here (starts with https:// or http://) or
+# path to a local file, but in this case it should be absolute.
+urls = [
+    # "https://iplists.firehol.org/files/firehol_level1.netset",
+    # "/local.file"
+
+]
+update-each = "24h"
+
 # statsd statistics integration.
 [stats.statsd]
 # enabled/disabled

internal/cli/run_proxy.go

@@ -86,15 +86,15 @@ func makeAntiReplayCache(conf *config.Config) mtglib.AntiReplayCache {
 	)
 }
 
-func makeIPBlocklist(conf *config.Config, logger mtglib.Logger, ntw mtglib.Network) (mtglib.IPBlocklist, error) {
-	if !conf.Defense.Blocklist.Enabled.Get(false) {
+func makeIPBlocklist(conf config.ListConfig, logger mtglib.Logger, ntw mtglib.Network) (mtglib.IPBlocklist, error) {
+	if !conf.Enabled.Get(false) {
 		return ipblocklist.NewNoop(), nil
 	}
 
 	remoteURLs := []string{}
 	localFiles := []string{}
 
-	for _, v := range conf.Defense.Blocklist.URLs {
+	for _, v := range conf.URLs {
 		if v.IsRemote() {
 			remoteURLs = append(remoteURLs, v.String())
 		} else {
@@ -104,7 +104,7 @@ func makeIPBlocklist(conf *config.Config, logger mtglib.Logger, ntw mtglib.Netwo
 
 	firehol, err := ipblocklist.NewFirehol(logger.Named("ipblockist"),
 		ntw,
-		conf.Defense.Blocklist.DownloadConcurrency.Get(1),
+		conf.DownloadConcurrency.Get(1),
 		remoteURLs,
 		localFiles)
 	if err != nil {
@@ -153,7 +153,7 @@ func makeEventStream(conf *config.Config, logger mtglib.Logger) (mtglib.EventStr
 	return events.NewNoopStream(), nil
 }
 
-func runProxy(conf *config.Config, version string) error {
+func runProxy(conf *config.Config, version string) error { // nolint: funlen
 	logger := makeLogger(conf)
 
 	logger.BindJSON("configuration", conf.String()).Debug("configuration")
@@ -163,11 +163,22 @@ func runProxy(conf *config.Config, version string) error {
 		return fmt.Errorf("cannot build network: %w", err)
 	}
 
-	blocklist, err := makeIPBlocklist(conf, logger, ntw)
+	blocklist, err := makeIPBlocklist(conf.Defense.Blocklist, logger, ntw)
 	if err != nil {
 		return fmt.Errorf("cannot build ip blocklist: %w", err)
 	}
 
+	var whitelist mtglib.IPBlocklist
+
+	if conf.Defense.Allowlist.Enabled.Get(false) {
+		whlist, err := makeIPBlocklist(conf.Defense.Allowlist, logger, ntw)
+		if err != nil {
+			return fmt.Errorf("cannot build ip blocklist: %w", err)
+		}
+
+		whitelist = whlist
+	}
+
 	eventStream, err := makeEventStream(conf, logger)
 	if err != nil {
 		return fmt.Errorf("cannot build event stream: %w", err)
@@ -178,6 +189,7 @@ func runProxy(conf *config.Config, version string) error {
 		Network:         ntw,
 		AntiReplayCache: makeAntiReplayCache(conf),
 		IPBlocklist:     blocklist,
+		IPWhitelist:     whitelist,
 		EventStream:     eventStream,
 
 		Secret:             conf.Secret,

internal/config/config.go

@@ -8,6 +8,18 @@ import (
 	"github.com/9seconds/mtg/v2/mtglib"
 )
 
+type Optional struct {
+	Enabled TypeBool `json:"enabled"`
+}
+
+type ListConfig struct {
+	Optional
+
+	DownloadConcurrency TypeConcurrency    `json:"downloadConcurrency"`
+	URLs                []TypeBlocklistURI `json:"urls"`
+	UpdateEach          TypeDuration       `json:"updateEach"`
+}
+
 type Config struct {
 	Debug                    TypeBool        `json:"debug"`
 	AllowFallbackOnUnknownDC TypeBool        `json:"allowFallbackOnUnknownDc"`
@@ -20,16 +32,13 @@ type Config struct {
 	Concurrency              TypeConcurrency `json:"concurrency"`
 	Defense                  struct {
 		AntiReplay struct {
-			Enabled   TypeBool      `json:"enabled"`
+			Optional
+
 			MaxSize   TypeBytes     `json:"maxSize"`
 			ErrorRate TypeErrorRate `json:"errorRate"`
 		} `json:"antiReplay"`
-		Blocklist struct {
-			Enabled             TypeBool           `json:"enabled"`
-			DownloadConcurrency TypeConcurrency    `json:"downloadConcurrency"`
-			URLs                []TypeBlocklistURI `json:"urls"`
-			UpdateEach          TypeDuration       `json:"updateEach"`
-		} `json:"blocklist"`
+		Blocklist ListConfig `json:"blocklist"`
+		Allowlist ListConfig `json:"allowlist"`
 	} `json:"defense"`
 	Network struct {
 		Timeout struct {
@@ -42,13 +51,15 @@ type Config struct {
 	} `json:"network"`
 	Stats struct {
 		StatsD struct {
-			Enabled      TypeBool            `json:"enabled"`
+			Optional
+
 			Address      TypeHostPort        `json:"address"`
 			MetricPrefix TypeMetricPrefix    `json:"metricPrefix"`
 			TagFormat    TypeStatsdTagFormat `json:"tagFormat"`
 		} `json:"statsd"`
 		Prometheus struct {
-			Enabled      TypeBool         `json:"enabled"`
+			Optional
+
 			BindTo       TypeHostPort     `json:"bindTo"`
 			HTTPPath     TypeHTTPPath     `json:"httpPath"`
 			MetricPrefix TypeMetricPrefix `json:"metricPrefix"`

internal/config/parse.go

@@ -30,6 +30,12 @@ type tomlConfig struct {
 			URLs                []string `toml:"urls" json:"urls,omitempty"`
 			UpdateEach          string   `toml:"update-each" json:"updateEach,omitempty"`
 		} `toml:"blocklist" json:"blocklist,omitempty"`
+		Allowlist struct {
+			Enabled             bool     `toml:"enabled" json:"enabled,omitempty"`
+			DownloadConcurrency uint     `toml:"download-concurrency" json:"downloadConcurrency,omitempty"`
+			URLs                []string `toml:"urls" json:"urls,omitempty"`
+			UpdateEach          string   `toml:"update-each" json:"updateEach,omitempty"`
+		} `toml:"allowlist" json:"allowlist,omitempty"`
 	} `toml:"defense" json:"defense,omitempty"`
 	Network struct {
 		Timeout struct {

ipblocklist/files/http.go

@@ -0,0 +1,63 @@
+package files
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+)
+
+type httpFile struct {
+	http *http.Client
+	url  string
+}
+
+func (h httpFile) Open(ctx context.Context) (io.ReadCloser, error) {
+	request, err := http.NewRequestWithContext(ctx, http.MethodGet, h.url, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	response, err := h.http.Do(request)
+	if err != nil {
+		if response != nil {
+			io.Copy(io.Discard, response.Body) // nolint: errcheck
+			response.Body.Close()
+		}
+
+		return nil, fmt.Errorf("cannot get url %s: %w", h.url, err)
+	}
+
+	if response.StatusCode >= http.StatusBadRequest {
+		return nil, fmt.Errorf("unexpected status code %d", response.StatusCode)
+	}
+
+	return response.Body, nil
+}
+
+func (h httpFile) String() string {
+	return h.url
+}
+
+func NewHTTP(client *http.Client, endpoint string) (File, error) {
+	if client == nil {
+		return nil, ErrBadHTTPClient
+	}
+
+	parsed, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, fmt.Errorf("incorrect url %s: %w", endpoint, err)
+	}
+
+	switch parsed.Scheme {
+	case "http", "https":
+	default:
+		return nil, fmt.Errorf("unsupported url %s", endpoint)
+	}
+
+	return httpFile{
+		http: client,
+		url:  endpoint,
+	}, nil
+}

ipblocklist/files/http_test.go

@@ -0,0 +1,90 @@
+package files_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
+	"github.com/stretchr/testify/suite"
+)
+
+type HTTPTestSuite struct {
+	suite.Suite
+
+	httpClient *http.Client
+	httpServer *httptest.Server
+	ctx        context.Context
+	ctxCancel  context.CancelFunc
+}
+
+func (suite *HTTPTestSuite) makeFile(path string) (files.File, error) {
+	return files.NewHTTP(suite.httpClient, suite.httpServer.URL+"/"+path) // nolint: wrapcheck
+}
+
+func (suite *HTTPTestSuite) SetupSuite() {
+	mux := http.NewServeMux()
+
+	mux.Handle("/", http.FileServer(http.Dir("testdata")))
+
+	suite.httpServer = httptest.NewServer(mux)
+	suite.httpClient = suite.httpServer.Client()
+}
+
+func (suite *HTTPTestSuite) SetupTest() {
+	suite.ctx, suite.ctxCancel = context.WithCancel(context.Background())
+}
+
+func (suite *HTTPTestSuite) TearDownTest() {
+	suite.ctxCancel()
+	suite.httpServer.CloseClientConnections()
+}
+
+func (suite *HTTPTestSuite) TearDownSuite() {
+	suite.httpServer.Close()
+}
+
+func (suite *HTTPTestSuite) TestBadURL() {
+	_, err := files.NewHTTP(suite.httpClient, "sdfsdf")
+	suite.Error(err)
+}
+
+func (suite *HTTPTestSuite) TestBadSchema() {
+	_, err := files.NewHTTP(suite.httpClient, "gopher://lala")
+	suite.Error(err)
+}
+
+func (suite *HTTPTestSuite) TestNilHTTPClient() {
+	_, err := files.NewHTTP(nil, "")
+	suite.Error(err)
+}
+
+func (suite *HTTPTestSuite) TestAbsentFile() {
+	file, err := suite.makeFile("absent")
+	suite.NoError(err)
+
+	_, err = file.Open(suite.ctx)
+	suite.Error(err)
+}
+
+func (suite *HTTPTestSuite) TestOk() {
+	file, err := suite.makeFile("readable")
+	suite.NoError(err)
+
+	readCloser, err := file.Open(suite.ctx)
+	suite.NoError(err)
+
+	defer readCloser.Close()
+
+	data, err := io.ReadAll(readCloser)
+	suite.NoError(err)
+	suite.Equal("Hooray!", strings.TrimSpace(string(data)))
+}
+
+func TestHTTP(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &HTTPTestSuite{})
+}

ipblocklist/files/init.go

@@ -0,0 +1,14 @@
+package files
+
+import (
+	"context"
+	"errors"
+	"io"
+)
+
+var ErrBadHTTPClient = errors.New("incorrect http client")
+
+type File interface {
+	Open(context.Context) (io.ReadCloser, error)
+	String() string
+}

ipblocklist/files/local.go

@@ -0,0 +1,30 @@
+package files
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+)
+
+type localFile struct {
+	path string
+}
+
+func (l localFile) Open(ctx context.Context) (io.ReadCloser, error) {
+	return os.Open(l.path) // nolint: wrapcheck
+}
+
+func (l localFile) String() string {
+	return l.path
+}
+
+func NewLocal(path string) (File, error) {
+	if stat, err := os.Stat(path); os.IsNotExist(err) || stat.IsDir() || stat.Mode().Perm()&0o400 == 0 {
+		return nil, fmt.Errorf("%s is not a readable file", path)
+	}
+
+	return localFile{
+		path: path,
+	}, nil
+}

ipblocklist/files/local_test.go

@@ -0,0 +1,55 @@
+package files_test
+
+import (
+	"context"
+	"io"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
+)
+
+type LocalTestSuite struct {
+	suite.Suite
+}
+
+func (suite *LocalTestSuite) getLocalFile(name string) string {
+	return filepath.Join("testdata", name)
+}
+
+func (suite *LocalTestSuite) TestIncorrect() {
+	names := []string{
+		"absent",
+		"directory",
+	}
+
+	for _, v := range names {
+		value := v
+
+		suite.T().Run(v, func(t *testing.T) {
+			_, err := files.NewLocal(suite.getLocalFile(value))
+			assert.Error(t, err)
+		})
+	}
+}
+
+func (suite *LocalTestSuite) TestOk() {
+	file, err := files.NewLocal(suite.getLocalFile("readable"))
+	suite.NoError(err)
+
+	reader, err := file.Open(context.Background())
+	suite.NoError(err)
+
+	data, err := io.ReadAll(reader)
+	suite.NoError(err)
+
+	suite.Equal("Hooray!", strings.TrimSpace(string(data)))
+}
+
+func TestLocal(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &LocalTestSuite{})
+}

ipblocklist/files/testdata/readable

@@ -0,0 +1 @@
+Hooray!

ipblocklist/firehol.go

@@ -4,16 +4,13 @@ import (
 	"bufio"
 	"context"
 	"fmt"
-	"io"
 	"net"
-	"net/http"
-	"net/url"
-	"os"
 	"regexp"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
 	"github.com/9seconds/mtg/v2/mtglib"
 	"github.com/kentik/patricia"
 	"github.com/kentik/patricia/bool_tree"
@@ -41,20 +38,16 @@ var fireholRegexpComment = regexp.MustCompile(`\s*#.*?$`)
 //     127.0.0.1   # you can specify an IP
 //     10.0.0.0/8  # or cidr
 type Firehol struct {
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-	logger    mtglib.Logger
+	ctx         context.Context
+	ctxCancel   context.CancelFunc
+	logger      mtglib.Logger
+	updateMutex sync.RWMutex
 
-	rwMutex sync.RWMutex
+	blocklists []files.File
 
-	remoteURLs []string
-	localFiles []string
-
-	httpClient *http.Client
 	workerPool *ants.Pool
-
-	treeV4 *bool_tree.TreeV4
-	treeV6 *bool_tree.TreeV6
+	treeV4     *bool_tree.TreeV4
+	treeV6     *bool_tree.TreeV6
 }
 
 // Shutdown stop a background update process.
@@ -68,8 +61,8 @@ func (f *Firehol) Contains(ip net.IP) bool {
 		return true
 	}
 
-	f.rwMutex.RLock()
-	defer f.rwMutex.RUnlock()
+	f.updateMutex.RLock()
+	defer f.updateMutex.RUnlock()
 
 	if ip4 := ip.To4(); ip4 != nil {
 		return f.containsIPv4(ip4)
@@ -98,22 +91,14 @@ func (f *Firehol) Run(updateEach time.Duration) {
 		}
 	}()
 
-	if err := f.update(); err != nil {
-		f.logger.WarningError("cannot update blocklist", err)
-	} else {
-		f.logger.Info("blocklist was updated")
-	}
+	f.update()
 
 	for {
 		select {
 		case <-f.ctx.Done():
 			return
 		case <-ticker.C:
-			if err := f.update(); err != nil {
-				f.logger.WarningError("cannot update blocklist", err)
-			} else {
-				f.logger.Info("blocklist was updated")
-			}
+			f.update()
 		}
 	}
 }
@@ -138,121 +123,53 @@ func (f *Firehol) containsIPv6(addr net.IP) bool {
 	return false
 }
 
-func (f *Firehol) update() error { // nolint: funlen, cyclop
+func (f *Firehol) update() {
 	ctx, cancel := context.WithCancel(f.ctx)
 	defer cancel()
 
 	wg := &sync.WaitGroup{}
-	wg.Add(len(f.remoteURLs) + len(f.localFiles))
+	wg.Add(len(f.blocklists))
 
 	treeMutex := &sync.Mutex{}
 	v4tree := bool_tree.NewTreeV4()
 	v6tree := bool_tree.NewTreeV6()
 
-	errorChan := make(chan error, 1)
-	defer close(errorChan)
-
-	for _, v := range f.localFiles {
-		go func(filename string) {
+	for _, v := range f.blocklists {
+		go func(file files.File) {
 			defer wg.Done()
 
-			if err := f.updateLocalFile(ctx, filename, treeMutex, v4tree, v6tree); err != nil {
-				cancel()
-				f.logger.BindStr("filename", filename).WarningError("cannot update", err)
+			logger := f.logger.BindStr("filename", file.String())
 
-				select {
-				case errorChan <- err:
-				default:
-				}
-			}
-		}(v)
-	}
+			fileContent, err := file.Open(ctx)
+			if err != nil {
+				logger.WarningError("update has failed", err)
 
-	for _, v := range f.remoteURLs {
-		value := v
-
-		f.workerPool.Submit(func() { // nolint: errcheck
-			defer wg.Done()
+				return
+			}
 
-			if err := f.updateRemoteURL(ctx, value, treeMutex, v4tree, v6tree); err != nil {
-				cancel()
-				f.logger.BindStr("url", value).WarningError("cannot update", err)
+			defer fileContent.Close()
 
-				select {
-				case errorChan <- err:
-				default:
-				}
+			if err := f.updateFromFile(treeMutex, v4tree, v6tree, bufio.NewScanner(fileContent)); err != nil {
+				logger.WarningError("update has failed", err)
 			}
-		})
+		}(v)
 	}
 
 	wg.Wait()
 
-	select {
-	case err := <-errorChan:
-		return fmt.Errorf("cannot update trees: %w", err)
-	default:
-	}
-
-	f.rwMutex.Lock()
-	defer f.rwMutex.Unlock()
+	f.updateMutex.Lock()
+	defer f.updateMutex.Unlock()
 
 	f.treeV4 = v4tree
 	f.treeV6 = v6tree
 
-	return nil
-}
-
-func (f *Firehol) updateLocalFile(ctx context.Context, filename string,
-	mutex sync.Locker,
-	v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
-	filefp, err := os.Open(filename)
-	if err != nil {
-		return fmt.Errorf("cannot open file: %w", err)
-	}
-
-	go func(ctx context.Context, closer io.Closer) {
-		<-ctx.Done()
-		closer.Close()
-	}(ctx, filefp)
-
-	defer filefp.Close()
-
-	return f.updateTrees(mutex, filefp, v4tree, v6tree)
+	f.logger.Info("blocklist was updated")
 }
 
-func (f *Firehol) updateRemoteURL(ctx context.Context, url string,
-	mutex sync.Locker,
-	v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
-	if err != nil {
-		return fmt.Errorf("cannot build a request: %w", err)
-	}
-
-	resp, err := f.httpClient.Do(req) // nolint: bodyclose
-	if err != nil {
-		return fmt.Errorf("cannot request a remote URL %s: %w", url, err)
-	}
-
-	go func(ctx context.Context, closer io.Closer) {
-		<-ctx.Done()
-		closer.Close()
-	}(ctx, resp.Body)
-
-	defer func(rc io.ReadCloser) {
-		io.Copy(io.Discard, rc) // nolint: errcheck
-		rc.Close()
-	}(resp.Body)
-
-	return f.updateTrees(mutex, resp.Body, v4tree, v6tree)
-}
-
-func (f *Firehol) updateTrees(mutex sync.Locker,
-	reader io.Reader,
+func (f *Firehol) updateFromFile(mutex sync.Locker,
 	v4tree *bool_tree.TreeV4,
-	v6tree *bool_tree.TreeV6) error {
-	scanner := bufio.NewScanner(reader)
-
+	v6tree *bool_tree.TreeV6,
+	scanner *bufio.Scanner) error {
 	for scanner.Scan() {
 		text := scanner.Text()
 		text = fireholRegexpComment.ReplaceAllLiteralString(text, "")
@@ -271,7 +188,7 @@ func (f *Firehol) updateTrees(mutex sync.Locker,
 	}
 
 	if scanner.Err() != nil {
-		return fmt.Errorf("cannot parse a response: %w", scanner.Err())
+		return fmt.Errorf("cannot parse a file: %w", scanner.Err())
 	}
 
 	return nil
@@ -317,27 +234,36 @@ func (f *Firehol) updateAddToTrees(ip net.IP, cidr uint,
 // when it is necessary.
 func NewFirehol(logger mtglib.Logger, network mtglib.Network,
 	downloadConcurrency uint,
-	remoteURLs []string,
+	urls []string,
 	localFiles []string) (*Firehol, error) {
-	for _, v := range remoteURLs {
-		parsed, err := url.Parse(v)
+	blocklists := []files.File{}
+
+	for _, v := range localFiles {
+		file, err := files.NewLocal(v)
 		if err != nil {
-			return nil, fmt.Errorf("incorrect url %s: %w", v, err)
+			return nil, fmt.Errorf("cannot create a local file %s: %w", v, err)
 		}
 
-		switch parsed.Scheme {
-		case "http", "https":
-		default:
-			return nil, fmt.Errorf("unsupported url %s", v)
-		}
+		blocklists = append(blocklists, file)
 	}
 
-	for _, v := range localFiles {
-		if stat, err := os.Stat(v); os.IsNotExist(err) || stat.IsDir() || stat.Mode().Perm()&0o400 == 0 {
-			return nil, fmt.Errorf("%s is not a readable file", v)
+	httpClient := network.MakeHTTPClient(nil)
+
+	for _, v := range urls {
+		file, err := files.NewHTTP(httpClient, v)
+		if err != nil {
+			return nil, fmt.Errorf("cannot create a HTTP file %s: %w", v, err)
 		}
+
+		blocklists = append(blocklists, file)
 	}
 
+	return NewFireholFromFiles(logger, downloadConcurrency, blocklists)
+}
+
+func NewFireholFromFiles(logger mtglib.Logger,
+	downloadConcurrency uint,
+	blocklists []files.File) (*Firehol, error) {
 	if downloadConcurrency == 0 {
 		downloadConcurrency = DefaultFireholDownloadConcurrency
 	}
@@ -349,11 +275,9 @@ func NewFirehol(logger mtglib.Logger, network mtglib.Network,
 		ctx:        ctx,
 		ctxCancel:  cancel,
 		logger:     logger.Named("firehol"),
-		httpClient: network.MakeHTTPClient(nil),
 		treeV4:     bool_tree.NewTreeV4(),
 		treeV6:     bool_tree.NewTreeV6(),
 		workerPool: workerPool,
-		remoteURLs: remoteURLs,
-		localFiles: localFiles,
+		blocklists: blocklists,
 	}, nil
 }

mtglib/proxy.go

@@ -33,7 +33,8 @@ type Proxy struct {
 	secret          Secret
 	network         Network
 	antiReplayCache AntiReplayCache
-	ipBlocklist     IPBlocklist
+	blocklist       IPBlocklist
+	whitelist       IPBlocklist
 	eventStream     EventStream
 	logger          Logger
 }
@@ -91,7 +92,7 @@ func (p *Proxy) ServeConn(conn net.Conn) {
 }
 
 // Serve starts a proxy on a given listener.
-func (p *Proxy) Serve(listener net.Listener) error {
+func (p *Proxy) Serve(listener net.Listener) error { // nolint: cyclop
 	p.streamWaitGroup.Add(1)
 	defer p.streamWaitGroup.Done()
 
@@ -109,7 +110,15 @@ func (p *Proxy) Serve(listener net.Listener) error {
 		ipAddr := conn.RemoteAddr().(*net.TCPAddr).IP
 		logger := p.logger.BindStr("ip", ipAddr.String())
 
-		if p.ipBlocklist.Contains(ipAddr) {
+		if p.whitelist != nil && !p.whitelist.Contains(ipAddr) {
+			conn.Close()
+			logger.Info("ip was rejected by whitelist")
+			p.eventStream.Send(p.ctx, NewEventIPBlocklisted(ipAddr))
+
+			continue
+		}
+
+		if p.blocklist.Contains(ipAddr) {
 			conn.Close()
 			logger.Info("ip was blacklisted")
 			p.eventStream.Send(p.ctx, NewEventIPBlocklisted(ipAddr))
@@ -291,7 +300,8 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		secret:                   opts.Secret,
 		network:                  opts.Network,
 		antiReplayCache:          opts.AntiReplayCache,
-		ipBlocklist:              opts.IPBlocklist,
+		blocklist:                opts.IPBlocklist,
+		whitelist:                opts.IPWhitelist,
 		eventStream:              opts.EventStream,
 		logger:                   opts.getLogger("proxy"),
 		domainFrontingPort:       opts.getDomainFrontingPort(),

mtglib/proxy_opts.go

@@ -28,6 +28,11 @@ type ProxyOpts struct {
 	// This is a mandatory setting.
 	IPBlocklist IPBlocklist
 
+	// IPWhitelist defines a whitelist of IPs to allow to use proxy.
+	//
+	// This is an optional setting, ignored by default (no restrictions).
+	IPWhitelist IPBlocklist
+
 	// EventStream defines an instance of event stream.
 	//
 	// This ia a mandatory setting.