diff --git a/go.mod b/go.mod index 05e2b83..a09b213 100644 --- a/go.mod +++ b/go.mod @@ -3,26 +3,25 @@ module github.com/9seconds/mtg/v2 go 1.16 require ( - github.com/OneOfOne/xxhash v1.2.8 // indirect + github.com/OneOfOne/xxhash v1.2.8 github.com/alecthomas/kong v0.2.16 github.com/alecthomas/units v0.0.0-20210208195552-ff826a37aa15 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 github.com/babolivier/go-doh-client v0.0.0-20201028162107-a76cff4cb8b6 + github.com/d4l3k/messagediff v1.2.1 // indirect github.com/jarcoal/httpmock v1.0.8 - github.com/kentik/patricia v0.0.0-20201202224819-f9447a6e25f1 // indirect - github.com/kr/pretty v0.1.0 // indirect + github.com/kentik/patricia v0.0.0-20201202224819-f9447a6e25f1 github.com/libp2p/go-reuseport v0.0.2 github.com/mccutchen/go-httpbin v1.1.1 - github.com/panjf2000/ants v1.3.0 // indirect + github.com/panjf2000/ants/v2 v2.4.3 github.com/pelletier/go-toml v1.8.1 - github.com/prometheus/client_golang v1.9.0 // indirect - github.com/rs/zerolog v1.20.0 // indirect - github.com/smira/go-statsd v1.3.2 // indirect + github.com/prometheus/client_golang v1.9.0 + github.com/rs/zerolog v1.20.0 + github.com/smira/go-statsd v1.3.2 github.com/stretchr/objx v0.3.0 // indirect github.com/stretchr/testify v1.7.0 - github.com/tylertreat/BoomFilters v0.0.0-20200520150052-42a7b4300c0c // indirect + github.com/tylertreat/BoomFilters v0.0.0-20200520150052-42a7b4300c0c github.com/xeipuuv/gojsonschema v1.2.0 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110 golang.org/x/sys v0.0.0-20210309074719-68d13333faf2 // indirect - gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect ) diff --git a/go.sum b/go.sum index fd76767..d9e0113 100644 --- a/go.sum +++ b/go.sum @@ -50,6 +50,8 @@ github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7 github.com/coreos/pkg v0.0.0-20160727233714-3ac0863d7acf/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA= github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY= +github.com/d4l3k/messagediff v1.2.1 h1:ZcAIMYsUg0EAp9X+tt8/enBE/Q8Yd5kzPynLyKptt9U= +github.com/d4l3k/messagediff v1.2.1/go.mod h1:Oozbb1TVXFac9FtSIxHBMnBCq2qeH/2KkEQxENCrlLo= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -100,6 +102,7 @@ github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= @@ -210,8 +213,8 @@ github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJ github.com/openzipkin/zipkin-go v0.2.1/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4= github.com/openzipkin/zipkin-go v0.2.2/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4= github.com/pact-foundation/pact-go v1.0.4/go.mod h1:uExwJY4kCzNPcHRj+hCR/HBbOOIwwtUjcrb0b5/5kLM= -github.com/panjf2000/ants v1.3.0 h1:8pQ+8leaLc9lys2viEEr8md0U4RN6uOSUCE9bOYjQ9M= -github.com/panjf2000/ants v1.3.0/go.mod h1:AaACblRPzq35m1g3enqYcxspbbiOJJYaxU2wMpm1cXY= +github.com/panjf2000/ants/v2 v2.4.3 h1:wHghL17YKFanB62QjPQ9o+DuM4q7WrQ7zAhoX8+eBXU= +github.com/panjf2000/ants/v2 v2.4.3/go.mod h1:f6F0NZVFsGCp5A7QW/Zj/m92atWwOkY0OIhFxRNFr4A= github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc= github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k= github.com/pelletier/go-toml v1.8.1 h1:1Nf83orprkJyknT6h7zbuEGUEjcyVlCxSUGTENmNCRM= @@ -402,6 +405,7 @@ golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.0.0-20200103221440-774c71fcf114/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/api v0.3.1/go.mod h1:6wY9I6uQWHQ8EM57III9mq/AjF+i8G65rmVagqKMtkk= google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= @@ -445,6 +449,7 @@ gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.5/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/ipblocklist/firehol.go b/ipblocklist/firehol.go index 876ad82..2d8a748 100644 --- a/ipblocklist/firehol.go +++ b/ipblocklist/firehol.go @@ -18,7 +18,7 @@ import ( "github.com/9seconds/mtg/v2/mtglib" "github.com/kentik/patricia" "github.com/kentik/patricia/bool_tree" - "github.com/panjf2000/ants" + "github.com/panjf2000/ants/v2" ) const ( diff --git a/mtglib/init.go b/mtglib/init.go index 22c5472..e663c62 100644 --- a/mtglib/init.go +++ b/mtglib/init.go @@ -8,7 +8,19 @@ import ( "time" ) -var ErrSecretEmpty = errors.New("secret is empty") +var ( + ErrSecretEmpty = errors.New("secret is empty") + ErrSecretInvalid = errors.New("secret is invalid") + ErrNetworkIsNotDefined = errors.New("network is not defined") + ErrAntiReplayCacheIsNotDefined = errors.New("anti-replay cache is not defined") + ErrIPBlocklistIsNotDefined = errors.New("ip blocklist is not defined") + ErrEventStreamIsNotDefined = errors.New("event stream is not defined") + ErrLoggerIsNotDefined = errors.New("logger is not defined") +) + +const ( + DefaultConcurrency = 4096 +) type Network interface { Dial(network, address string) (net.Conn, error) diff --git a/mtglib/proxy.go b/mtglib/proxy.go new file mode 100644 index 0000000..e603db1 --- /dev/null +++ b/mtglib/proxy.go @@ -0,0 +1,120 @@ +package mtglib + +import ( + "context" + "errors" + "fmt" + "net" + "sync" + "time" + + "github.com/panjf2000/ants/v2" +) + +type Proxy struct { + ctx context.Context + ctxCancel context.CancelFunc + streamWaitGroup sync.WaitGroup + workerPool *ants.PoolWithFunc + + secret Secret + network Network + antiReplayCache AntiReplayCache + ipBlocklist IPBlocklist + eventStream EventStream + logger Logger +} + +func (p *Proxy) ServeConn(conn net.Conn) { + ctx := newStreamContext(p.ctx, p.logger, conn) + defer ctx.Close() + + p.eventStream.Send(ctx, EventStart{ + CreatedAt: time.Now(), + ConnID: ctx.connID, + RemoteIP: ctx.ClientIP(), + }) + ctx.logger.Info("Stream has been started") + + defer func() { + p.eventStream.Send(ctx, EventFinish{ + CreatedAt: time.Now(), + ConnID: ctx.connID, + }) + ctx.logger.Info("Stream has been finished") + }() +} + +func (p *Proxy) Serve(listener net.Listener) error { + for { + conn, err := listener.Accept() + if err != nil { + return fmt.Errorf("cannot accept a new connection: %w", err) + } + + err = p.workerPool.Invoke(conn) + + switch { + case err == nil: + case errors.Is(err, ants.ErrPoolClosed): + return nil + case errors.Is(err, ants.ErrPoolOverload): + p.eventStream.Send(p.ctx, EventConcurrencyLimited{}) + } + } +} + +func (p *Proxy) Shutdown() { + p.ctxCancel() + p.streamWaitGroup.Wait() + p.workerPool.Release() +} + +type antsLogger struct{} + +func (a antsLogger) Printf(msg string, args ...interface{}) {} + +func NewProxy(opts ProxyOpts) (*Proxy, error) { + switch { + case opts.Network == nil: + return nil, ErrNetworkIsNotDefined + case opts.AntiReplayCache == nil: + return nil, ErrAntiReplayCacheIsNotDefined + case opts.IPBlocklist == nil: + return nil, ErrIPBlocklistIsNotDefined + case opts.EventStream == nil: + return nil, ErrEventStreamIsNotDefined + case opts.Logger == nil: + return nil, ErrLoggerIsNotDefined + case !opts.Secret.Valid(): + return nil, ErrSecretInvalid + } + + concurrency := opts.Concurrency + if concurrency == 0 { + concurrency = DefaultConcurrency + } + + ctx, cancel := context.WithCancel(context.Background()) + proxy := &Proxy{ + ctx: ctx, + ctxCancel: cancel, + secret: opts.Secret, + network: opts.Network, + antiReplayCache: opts.AntiReplayCache, + ipBlocklist: opts.IPBlocklist, + eventStream: opts.EventStream, + logger: opts.Logger.Named("proxy"), + } + + pool, err := ants.NewPoolWithFunc(int(concurrency), func(arg interface{}) { + proxy.ServeConn(arg.(net.Conn)) + }, ants.WithLogger(antsLogger{})) + if err != nil { + return nil, fmt.Errorf("cannot initialize a pool: %w", err) + } + + proxy.workerPool = pool + + return proxy, nil +} diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go new file mode 100644 index 0000000..88130ca --- /dev/null +++ b/mtglib/proxy_opts.go @@ -0,0 +1,12 @@ +package mtglib + +type ProxyOpts struct { + Secret Secret + Network Network + AntiReplayCache AntiReplayCache + IPBlocklist IPBlocklist + EventStream EventStream + Logger Logger + + Concurrency uint +} diff --git a/mtglib/stream_context.go b/mtglib/stream_context.go new file mode 100644 index 0000000..cd5d1e8 --- /dev/null +++ b/mtglib/stream_context.go @@ -0,0 +1,63 @@ +package mtglib + +import ( + "context" + "crypto/rand" + "encoding/base64" + "net" + "time" +) + +type streamContext struct { + ctx context.Context + ctxCancel context.CancelFunc + clientConn net.Conn + connID string + logger Logger +} + +func (s *streamContext) Deadline() (time.Time, bool) { + return s.ctx.Deadline() +} + +func (s *streamContext) Done() <-chan struct{} { + return s.ctx.Done() +} + +func (s *streamContext) Err() error { + return s.ctx.Err() +} + +func (s *streamContext) Value(key interface{}) interface{} { + return s.ctx.Value(key) +} + +func (s *streamContext) Close() { + s.ctxCancel() + s.clientConn.Close() +} + +func (s *streamContext) ClientIP() net.IP { + return s.clientConn.RemoteAddr().(*net.TCPAddr).IP +} + +func newStreamContext(ctx context.Context, logger Logger, clientConn net.Conn) *streamContext { + connIDBytes := make([]byte, 16) + + if _, err := rand.Read(connIDBytes); err != nil { + panic(err) + } + + ctx, cancel := context.WithCancel(ctx) + streamCtx := &streamContext{ + ctx: ctx, + ctxCancel: cancel, + clientConn: clientConn, + connID: base64.RawURLEncoding.EncodeToString(connIDBytes), + } + streamCtx.logger = logger. + BindStr("stream-id", streamCtx.connID). + BindStr("client-ip", streamCtx.ClientIP().String()) + + return streamCtx +}