ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 17:04:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #262 from 9seconds/better-whitelist

Browse Source 
More elegant management of ip allowlists
This commit is contained in:
Sergey Arkhipov
2022-03-21 12:13:25 +03:00
committed by GitHub
parent 78a9ff18a5 a27facaa16
commit eba3673e27
7 changed files with 169 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/cli/run_proxy.go
+51 -19
View File
@@ -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,
ipblocklist/files/mem.go
+37
View File
@@ -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(),
}
}
ipblocklist/files/mem_test.go
+42
View File
@@ -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{})
}
mtglib/init.go
+4
View File
@@ -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
+6 -9
View File
@@ -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
+4 -2
View File
@@ -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
+25
View File
@@ -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