More elegant management of ip allowlists

internal/cli/run_proxy.go

@@ -12,11 +12,13 @@ import (
 	"github.com/9seconds/mtg/v2/internal/config"
 	"github.com/9seconds/mtg/v2/internal/utils"
 	"github.com/9seconds/mtg/v2/ipblocklist"
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
 	"github.com/9seconds/mtg/v2/logger"
 	"github.com/9seconds/mtg/v2/mtglib"
 	"github.com/9seconds/mtg/v2/network"
 	"github.com/9seconds/mtg/v2/stats"
 	"github.com/rs/zerolog"
+	"github.com/yl2chen/cidranger"
 )
 
 func makeLogger(conf *config.Config) mtglib.Logger {
@@ -106,7 +108,7 @@ func makeIPBlocklist(conf config.ListConfig,
 		}
 	}
 
-	firehol, err := ipblocklist.NewFirehol(logger.Named("ipblockist"),
+	blocklist, err := ipblocklist.NewFirehol(logger.Named("ipblockist"),
 		ntw,
 		conf.DownloadConcurrency.Get(1),
 		remoteURLs,
@@ -116,9 +118,44 @@ func makeIPBlocklist(conf config.ListConfig,
 		return nil, fmt.Errorf("incorrect parameters for firehol: %w", err)
 	}
 
-	go firehol.Run(conf.UpdateEach.Get(ipblocklist.DefaultFireholUpdateEach))
+	go blocklist.Run(conf.UpdateEach.Get(ipblocklist.DefaultFireholUpdateEach))
 
-	return firehol, nil
+	return blocklist, nil
+}
+
+func makeIPAllowlist(conf config.ListConfig,
+	logger mtglib.Logger,
+	ntw mtglib.Network,
+	updateCallback ipblocklist.FireholUpdateCallback,
+) (allowlist mtglib.IPBlocklist, err error) {
+	if !conf.Enabled.Get(false) {
+		allowlist, err = ipblocklist.NewFireholFromFiles(
+			logger.Named("ipblocklist"),
+			1,
+			[]files.File{
+				files.NewMem([]*net.IPNet{
+					cidranger.AllIPv4,
+					cidranger.AllIPv6,
+				}),
+			},
+			updateCallback,
+		)
+
+		go allowlist.Run(conf.UpdateEach.Get(ipblocklist.DefaultFireholUpdateEach))
+	} else {
+		allowlist, err = makeIPBlocklist(
+			conf,
+			logger,
+			ntw,
+			updateCallback,
+		)
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("cannot build allowlist: %w", err)
+	}
+
+	return allowlist, nil
 }
 
 func makeEventStream(conf *config.Config, logger mtglib.Logger) (mtglib.EventStream, error) {
@@ -186,29 +223,24 @@ func runProxy(conf *config.Config, version string) error { // nolint: funlen
 		return fmt.Errorf("cannot build ip blocklist: %w", err)
 	}
 
-	var whitelist mtglib.IPBlocklist
+	allowlist, err := makeIPAllowlist(
-
-	if conf.Defense.Allowlist.Enabled.Get(false) {
-		whlist, err := makeIPBlocklist(
 		conf.Defense.Allowlist,
 		logger.Named("allowlist"),
 		ntw,
 		func(ctx context.Context, size int) {
 			eventStream.Send(ctx, mtglib.NewEventIPListSize(size, false))
-			})
+		},
+	)
 	if err != nil {
 		return fmt.Errorf("cannot build ip allowlist: %w", err)
 	}
 
-		whitelist = whlist
-	}
-
 	opts := mtglib.ProxyOpts{
 		Logger:          logger,
 		Network:         ntw,
 		AntiReplayCache: makeAntiReplayCache(conf),
 		IPBlocklist:     blocklist,
-		IPWhitelist:     whitelist,
+		IPAllowlist:     allowlist,
 		EventStream:     eventStream,
 
 		Secret:             conf.Secret,

ipblocklist/files/mem.go

@@ -19,7 +19,7 @@ func (m memFile) String() string {
 	return "mem"
 }
 
-func NewMem(networks []*net.IPNet) (File) {
+func NewMem(networks []*net.IPNet) File {
 	builder := strings.Builder{}
 
 	if len(networks) > 0 {

mtglib/init.go

@@ -48,6 +48,10 @@ var (
 	// create a proxy but ip blocklist instance is not defined.
 	ErrIPBlocklistIsNotDefined = errors.New("ip blocklist is not defined")
 
+	// ErrIPAllowlistIsNotDefined is returned if you are trying to
+	// create a proxy but ip allowlist instance is not defined.
+	ErrIPAllowlistIsNotDefined = errors.New("ip allowlist is not defined")
+
 	// ErrEventStreamIsNotDefined is returned if you are trying to create a
 	// proxy but event stream instance is not defined.
 	ErrEventStreamIsNotDefined = errors.New("event stream is not defined")

mtglib/proxy.go

@@ -34,7 +34,7 @@ type Proxy struct {
 	network         Network
 	antiReplayCache AntiReplayCache
 	blocklist       IPBlocklist
-	whitelist       IPBlocklist
+	allowlist       IPBlocklist
 	eventStream     EventStream
 	logger          Logger
 }
@@ -91,7 +91,7 @@ func (p *Proxy) ServeConn(conn essentials.Conn) {
 }
 
 // Serve starts a proxy on a given listener.
-func (p *Proxy) Serve(listener net.Listener) error { // nolint: cyclop
+func (p *Proxy) Serve(listener net.Listener) error {
 	p.streamWaitGroup.Add(1)
 	defer p.streamWaitGroup.Done()
 
@@ -109,9 +109,9 @@ func (p *Proxy) Serve(listener net.Listener) error { // nolint: cyclop
 		ipAddr := conn.RemoteAddr().(*net.TCPAddr).IP // nolint: forcetypeassert
 		logger := p.logger.BindStr("ip", ipAddr.String())
 
-		if p.whitelist != nil && !p.whitelist.Contains(ipAddr) {
+		if !p.allowlist.Contains(ipAddr) {
 			conn.Close()
-			logger.Info("ip was rejected by whitelist")
+			logger.Info("ip was rejected by allowlist")
 			p.eventStream.Send(p.ctx, NewEventIPBlocklisted(ipAddr))
 
 			continue
@@ -145,10 +145,7 @@ func (p *Proxy) Shutdown() {
 	p.streamWaitGroup.Wait()
 	p.workerPool.Release()
 
-	if p.whitelist != nil {
+	p.allowlist.Shutdown()
-		p.whitelist.Shutdown()
-	}
-
 	p.blocklist.Shutdown()
 }
 
@@ -308,7 +305,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		network:                  opts.Network,
 		antiReplayCache:          opts.AntiReplayCache,
 		blocklist:                opts.IPBlocklist,
-		whitelist:                opts.IPWhitelist,
+		allowlist:                opts.IPAllowlist,
 		eventStream:              opts.EventStream,
 		logger:                   opts.getLogger("proxy"),
 		domainFrontingPort:       opts.getDomainFrontingPort(),

mtglib/proxy_opts.go

@@ -28,10 +28,10 @@ type ProxyOpts struct {
 	// This is a mandatory setting.
 	IPBlocklist IPBlocklist
 
-	// IPWhitelist defines a whitelist of IPs to allow to use proxy.
+	// IPAllowlist defines a whitelist of IPs to allow to use proxy.
 	//
 	// This is an optional setting, ignored by default (no restrictions).
-	IPWhitelist IPBlocklist
+	IPAllowlist IPBlocklist
 
 	// EventStream defines an instance of event stream.
 	//
@@ -125,6 +125,8 @@ func (p ProxyOpts) valid() error {
 		return ErrAntiReplayCacheIsNotDefined
 	case p.IPBlocklist == nil:
 		return ErrIPBlocklistIsNotDefined
+	case p.IPAllowlist == nil:
+		return ErrIPAllowlistIsNotDefined
 	case p.EventStream == nil:
 		return ErrEventStreamIsNotDefined
 	case p.Logger == nil:

mtglib/proxy_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/9seconds/mtg/v2/antireplay"
 	"github.com/9seconds/mtg/v2/events"
 	"github.com/9seconds/mtg/v2/ipblocklist"
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
 	"github.com/9seconds/mtg/v2/logger"
 	"github.com/9seconds/mtg/v2/mtglib"
 	"github.com/9seconds/mtg/v2/network"
@@ -22,6 +23,7 @@ import (
 	"github.com/gotd/td/telegram/dcs"
 	"github.com/gotd/td/tg"
 	"github.com/stretchr/testify/suite"
+	"github.com/yl2chen/cidranger"
 )
 
 type ProxyTestSuite struct {
@@ -49,11 +51,26 @@ func (suite *ProxyTestSuite) SetupSuite() {
 	ntw, err := network.NewNetwork(dialer, "mtgtest", "1.1.1.1", 0)
 	suite.NoError(err)
 
+	allowlist, _ := ipblocklist.NewFireholFromFiles(
+		logger.NewNoopLogger(),
+		1,
+		[]files.File{
+			files.NewMem([]*net.IPNet{
+				cidranger.AllIPv4,
+				cidranger.AllIPv6,
+			}),
+		},
+		nil,
+	)
+
+	allowlist.Run(time.Second)
+
 	suite.opts = &mtglib.ProxyOpts{
 		Secret:          mtglib.GenerateSecret("httpbin.org"),
 		Network:         ntw,
 		AntiReplayCache: antireplay.NewNoop(),
 		IPBlocklist:     ipblocklist.NewNoop(),
+		IPAllowlist:     allowlist,
 		EventStream:     events.NewNoopStream(),
 		Logger:          logger.NewNoopLogger(),
 		UseTestDCs:      true,
@@ -114,6 +131,14 @@ func (suite *ProxyTestSuite) TestCannotInitNoIPBlocklist() {
 	suite.Error(err)
 }
 
+func (suite *ProxyTestSuite) TestCannotInitNoIPAllowlist() {
+	opts := *suite.opts
+	opts.IPAllowlist = nil
+
+	_, err := mtglib.NewProxy(opts)
+	suite.Error(err)
+}
+
 func (suite *ProxyTestSuite) TestCannotInitNoEventStream() {
 	opts := *suite.opts
 	opts.EventStream = nil