Add 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) {
+func makeIPBlocklist(conf config.ListConfig, logger mtglib.Logger, ntw mtglib.Network) (mtglib.IPBlocklist, error) {
-	if !conf.Defense.Blocklist.Enabled.Get(false) {
+	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 {
+		Blocklist ListConfig `json:"blocklist"`
-			Enabled             TypeBool           `json:"enabled"`
+		Allowlist ListConfig `json:"allowlist"`
-			DownloadConcurrency TypeConcurrency    `json:"downloadConcurrency"`
-			URLs                []TypeBlocklistURI `json:"urls"`
-			UpdateEach          TypeDuration       `json:"updateEach"`
-		} `json:"blocklist"`
 	} `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 {

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.