Add whitelist support

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 {