Add skeleton of the proxy

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/kentik/patricia v0.0.0-20201202224819-f9447a6e25f1
-	github.com/kr/pretty v0.1.0 // indirect
 	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/prometheus/client_golang v1.9.0
-	github.com/rs/zerolog v1.20.0 // indirect
+	github.com/rs/zerolog v1.20.0
-	github.com/smira/go-statsd v1.3.2 // indirect
+	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
 )

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/v2 v2.4.3 h1:wHghL17YKFanB62QjPQ9o+DuM4q7WrQ7zAhoX8+eBXU=
-github.com/panjf2000/ants v1.3.0/go.mod h1:AaACblRPzq35m1g3enqYcxspbbiOJJYaxU2wMpm1cXY=
+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=

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 (

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)

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
+}

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
+}

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
+}