More elegant management of ip allowlists

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.whitelist.Shutdown()
-	}
-
+	p.allowlist.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