Add base tests for proxy

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

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