From c0899d0743fe1fb997d0b22a236bd1752de90e91 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 21 Mar 2022 10:18:50 +0300 Subject: [PATCH 1/2] Add in-memory file --- ipblocklist/files/mem.go | 37 ++++++++++++++++++++++++++++++ ipblocklist/files/mem_test.go | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 ipblocklist/files/mem.go create mode 100644 ipblocklist/files/mem_test.go diff --git a/ipblocklist/files/mem.go b/ipblocklist/files/mem.go new file mode 100644 index 0000000..8668612 --- /dev/null +++ b/ipblocklist/files/mem.go @@ -0,0 +1,37 @@ +package files + +import ( + "context" + "io" + "net" + "strings" +) + +type memFile struct { + data string +} + +func (m memFile) Open(ctx context.Context) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader(m.data)), nil +} + +func (m memFile) String() string { + return "mem" +} + +func NewMem(networks []*net.IPNet) (File) { + builder := strings.Builder{} + + if len(networks) > 0 { + builder.WriteString(networks[0].String()) + } + + for i := 1; i < len(networks); i++ { + builder.WriteString("\n") + builder.WriteString(networks[i].String()) + } + + return memFile{ + data: builder.String(), + } +} diff --git a/ipblocklist/files/mem_test.go b/ipblocklist/files/mem_test.go new file mode 100644 index 0000000..e3142b2 --- /dev/null +++ b/ipblocklist/files/mem_test.go @@ -0,0 +1,42 @@ +package files_test + +import ( + "context" + "io" + "net" + "strings" + "testing" + + "github.com/9seconds/mtg/v2/ipblocklist/files" + "github.com/stretchr/testify/suite" +) + +type MemTestSuite struct { + suite.Suite +} + +func (suite *MemTestSuite) TestOk() { + _, network1, _ := net.ParseCIDR("192.168.0.1/24") + _, network2, _ := net.ParseCIDR("2001:0db8:85a3:0000:0000:8a2e:0370:7334/36") + + file := files.NewMem([]*net.IPNet{ + network1, + network2, + }) + + reader, err := file.Open(context.Background()) + suite.NoError(err) + + data, err := io.ReadAll(reader) + suite.NoError(err) + + strData := strings.TrimSpace(string(data)) + + suite.Contains(strData, "192.168.0.0/24") + suite.Contains(strData, "2001:db8:8000::/36") +} + +func TestMem(t *testing.T) { + t.Parallel() + suite.Run(t, &MemTestSuite{}) +} From a27facaa1643c14b1416e75be1b4058bbfc2e5ce Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 21 Mar 2022 11:09:56 +0300 Subject: [PATCH 2/2] More elegant management of ip allowlists --- internal/cli/run_proxy.go | 70 ++++++++++++++++++++++++++++----------- ipblocklist/files/mem.go | 4 +-- mtglib/init.go | 4 +++ mtglib/proxy.go | 15 ++++----- mtglib/proxy_opts.go | 6 ++-- mtglib/proxy_test.go | 25 ++++++++++++++ 6 files changed, 92 insertions(+), 32 deletions(-) diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go index 01ec512..01c6240 100644 --- a/internal/cli/run_proxy.go +++ b/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,21 +223,16 @@ func runProxy(conf *config.Config, version string) error { // nolint: funlen 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.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 + allowlist, err := makeIPAllowlist( + 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) } opts := mtglib.ProxyOpts{ @@ -208,7 +240,7 @@ func runProxy(conf *config.Config, version string) error { // nolint: funlen Network: ntw, AntiReplayCache: makeAntiReplayCache(conf), IPBlocklist: blocklist, - IPWhitelist: whitelist, + IPAllowlist: allowlist, EventStream: eventStream, Secret: conf.Secret, diff --git a/ipblocklist/files/mem.go b/ipblocklist/files/mem.go index 8668612..db23bef 100644 --- a/ipblocklist/files/mem.go +++ b/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 { @@ -32,6 +32,6 @@ func NewMem(networks []*net.IPNet) (File) { } return memFile{ - data: builder.String(), + data: builder.String(), } } diff --git a/mtglib/init.go b/mtglib/init.go index a245935..0e00db6 100644 --- a/mtglib/init.go +++ b/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") diff --git a/mtglib/proxy.go b/mtglib/proxy.go index 86e882f..99fe486 100644 --- a/mtglib/proxy.go +++ b/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(), diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index faad210..c10463c 100644 --- a/mtglib/proxy_opts.go +++ b/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: diff --git a/mtglib/proxy_test.go b/mtglib/proxy_test.go index 59963c6..b446949 100644 --- a/mtglib/proxy_test.go +++ b/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