diff --git a/mtglib/proxy.go b/mtglib/proxy.go index de76a45..36de1a8 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -42,6 +42,9 @@ func (p *Proxy) DomainFrontingAddress() string { } func (p *Proxy) ServeConn(conn net.Conn) { + p.streamWaitGroup.Add(1) + defer p.streamWaitGroup.Done() + ctx := newStreamContext(p.ctx, p.logger, conn) defer ctx.Close() @@ -91,20 +94,24 @@ func (p *Proxy) ServeConn(conn net.Conn) { } func (p *Proxy) Serve(listener net.Listener) error { + p.streamWaitGroup.Add(1) + defer p.streamWaitGroup.Done() + for { conn, err := listener.Accept() if err != nil { return fmt.Errorf("cannot accept a new connection: %w", err) } - if addr := conn.RemoteAddr().(*net.TCPAddr).IP; p.ipBlocklist.Contains(addr) { + ipAddr := conn.RemoteAddr().(*net.TCPAddr).IP + logger := p.logger.BindStr("ip", ipAddr.String()) + + if p.ipBlocklist.Contains(ipAddr) { conn.Close() - p.logger. - BindStr("ip", conn.RemoteAddr().(*net.TCPAddr).IP.String()). - Info("ip was blacklisted") + logger.Info("ip was blacklisted") p.eventStream.Send(p.ctx, EventIPBlocklisted{ CreatedAt: time.Now(), - RemoteIP: addr, + RemoteIP: ipAddr, }) continue @@ -117,13 +124,17 @@ func (p *Proxy) Serve(listener net.Listener) error { case errors.Is(err, ants.ErrPoolClosed): return nil case errors.Is(err, ants.ErrPoolOverload): - p.logger. - BindStr("ip", conn.RemoteAddr().(*net.TCPAddr).IP.String()). - Info("connection was concurrency limited") + logger.Info("connection was concurrency limited") p.eventStream.Send(p.ctx, EventConcurrencyLimited{ CreatedAt: time.Now(), }) } + + select { + case <-p.ctx.Done(): + return p.ctx.Err() + default: + } } } @@ -292,9 +303,9 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop, funlen return nil, ErrSecretInvalid } - tg, err := telegram.New(opts.Network, opts.PreferIP) - if err != nil { - return nil, fmt.Errorf("cannot build telegram dialer: %w", err) + preferIP := opts.PreferIP + if preferIP == "" { + preferIP = DefaultPreferIP } concurrency := opts.Concurrency @@ -317,6 +328,11 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop, funlen domainFrontingPort = DefaultDomainFrontingPort } + tg, err := telegram.New(opts.Network, preferIP) + if err != nil { + return nil, fmt.Errorf("cannot build telegram dialer: %w", err) + } + ctx, cancel := context.WithCancel(context.Background()) proxy := &Proxy{ ctx: ctx, @@ -340,7 +356,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop, funlen ants.WithLogger(opts.Logger.Named("ants")), ants.WithNonblocking(true)) if err != nil { - return nil, fmt.Errorf("cannot initialize a pool: %w", err) + panic(err) } proxy.workerPool = pool diff --git a/mtglib/proxy_test.go b/mtglib/proxy_test.go new file mode 100644 index 0000000..6e767ee --- /dev/null +++ b/mtglib/proxy_test.go @@ -0,0 +1,184 @@ +package mtglib_test + +import ( + "crypto/tls" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "testing" + "time" + + "github.com/9seconds/mtg/v2/antireplay" + "github.com/9seconds/mtg/v2/events" + "github.com/9seconds/mtg/v2/ipblocklist" + "github.com/9seconds/mtg/v2/logger" + "github.com/9seconds/mtg/v2/mtglib" + "github.com/9seconds/mtg/v2/network" + "github.com/9seconds/mtg/v2/timeattack" + "github.com/stretchr/testify/suite" +) + +type ProxyTestSuite struct { + suite.Suite + + opts *mtglib.ProxyOpts + p *mtglib.Proxy + listener net.Listener +} + +func (suite *ProxyTestSuite) ProxyAddress() string { + _, port, _ := net.SplitHostPort(suite.listener.Addr().String()) + + return net.JoinHostPort("127.0.0.1", port) +} + +func (suite *ProxyTestSuite) ProxySecret() string { + return suite.opts.Secret.Hex() +} + +func (suite *ProxyTestSuite) SetupSuite() { + dialer, err := network.NewDefaultDialer(0, 0) + suite.NoError(err) + + ntw, err := network.NewNetwork(dialer, "mtgtest", "1.1.1.1", 0) + suite.NoError(err) + + suite.opts = &mtglib.ProxyOpts{ + Secret: mtglib.GenerateSecret("httpbin.org"), + Network: ntw, + AntiReplayCache: antireplay.NewNoop(), + TimeAttackDetector: timeattack.NewNoop(), + IPBlocklist: ipblocklist.NewNoop(), + EventStream: events.NewNoopStream(), + Logger: logger.NewNoopLogger(), + } + + proxy, err := mtglib.NewProxy(*suite.opts) + suite.NoError(err) + + suite.p = proxy + + listener, err := net.Listen("tcp", ":0") + suite.NoError(err) + + suite.listener = listener + + go suite.p.Serve(suite.listener) // nolint: errcheck +} + +func (suite *ProxyTestSuite) TearDownSuite() { + if suite.listener != nil { + suite.listener.Close() + } + + if suite.p != nil { + suite.p.Shutdown() + } +} + +func (suite *ProxyTestSuite) TestCannotInitNoSecret() { + opts := *suite.opts + opts.Secret = mtglib.Secret{} + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoNetwork() { + opts := *suite.opts + opts.Network = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoAntiReplayCache() { + opts := *suite.opts + opts.AntiReplayCache = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoIPBlocklist() { + opts := *suite.opts + opts.IPBlocklist = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoEventStream() { + opts := *suite.opts + opts.EventStream = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoTimeAttackDetector() { + opts := *suite.opts + opts.TimeAttackDetector = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitNoLogger() { + opts := *suite.opts + opts.Logger = nil + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestCannotInitIncorrectPreferIP() { + opts := *suite.opts + opts.PreferIP = "xxx" + + _, err := mtglib.NewProxy(opts) + suite.Error(err) +} + +func (suite *ProxyTestSuite) TestDomainFrontingAddress() { + suite.Equal("httpbin.org:443", suite.p.DomainFrontingAddress()) +} + +func (suite *ProxyTestSuite) TestHTTPSRequest() { + client := &http.Client{ + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{ + InsecureSkipVerify: true, + }, + }, + Timeout: 5 * time.Second, + } + + addr := fmt.Sprintf("https://%s/headers", suite.ProxyAddress()) + + resp, err := client.Get(addr) // nolint: noctx + suite.NoError(err) + + defer resp.Body.Close() + + suite.Equal(http.StatusOK, resp.StatusCode) + + data, err := io.ReadAll(resp.Body) + suite.NoError(err) + + jsonStruct := struct { + Headers struct { + TraceID string `json:"X-Amzn-Trace-Id"` + } `json:"headers"` + }{} + + suite.NoError(json.Unmarshal(data, &jsonStruct)) + suite.NotEmpty(jsonStruct.Headers.TraceID) +} + +func TestProxy(t *testing.T) { + t.Parallel() + suite.Run(t, &ProxyTestSuite{}) +}