mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 20:34:02 +03:00
More elegant management of ip allowlists
This commit is contained in:
+51
-19
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
+6
-9
@@ -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(),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user