Add support for domain fronting proxy protocol

mtglib/conns.go

@@ -3,9 +3,12 @@ package mtglib
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
+	"net"
 
 	"github.com/9seconds/mtg/v2/essentials"
+	"github.com/pires/go-proxyproto"
 )
 
 type connTraffic struct {
@@ -59,3 +62,36 @@ func newConnRewind(conn essentials.Conn) *connRewind {
 
 	return rv
 }
+
+type connProxyProtocol struct {
+	essentials.Conn
+
+	sourceAddr     net.Addr
+	headersWritten bool
+}
+
+func (c *connProxyProtocol) Write(p []byte) (int, error) {
+	if !c.headersWritten {
+		headers := proxyproto.HeaderProxyFromAddrs(2, c.sourceAddr, c.RemoteAddr())
+
+		toSend, err := headers.Format()
+		if err != nil {
+			panic(err)
+		}
+
+		if _, err := c.Conn.Write(toSend); err != nil {
+			return 0, fmt.Errorf("cannot send proxy protocol header: %w", err)
+		}
+
+		c.headersWritten = true
+	}
+
+	return c.Conn.Write(p)
+}
+
+func newConnProxyProtocol(source, target essentials.Conn) *connProxyProtocol {
+	return &connProxyProtocol{
+		Conn:       target,
+		sourceAddr: source.RemoteAddr(),
+	}
+}

mtglib/conns_internal_test.go

@@ -1,14 +1,17 @@
 package mtglib
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"errors"
 	"io"
+	"net"
 	"testing"
 	"time"
 
 	"github.com/9seconds/mtg/v2/internal/testlib"
+	"github.com/pires/go-proxyproto"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 )
@@ -200,6 +203,94 @@ func (suite *ConnRewindTestSuite) TestRead() {
 	suite.Equal([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, data)
 }
 
+type ConnProxyProtocolTestSuite struct {
+	suite.Suite
+
+	sourceConnMock *testlib.EssentialsConnMock
+	targetConnMock *testlib.EssentialsConnMock
+	conn           *connProxyProtocol
+}
+
+func (suite *ConnProxyProtocolTestSuite) SetupTest() {
+	suite.sourceConnMock = &testlib.EssentialsConnMock{}
+	suite.targetConnMock = &testlib.EssentialsConnMock{}
+
+	localAddr := &net.TCPAddr{
+		IP: net.ParseIP("127.0.0.1").To4(),
+	}
+	remoteAddr := &net.TCPAddr{
+		IP: net.ParseIP("127.0.0.2").To4(),
+	}
+
+	suite.sourceConnMock.
+		On("RemoteAddr").
+		Return(localAddr)
+	suite.targetConnMock.
+		On("RemoteAddr").
+		Maybe().
+		Return(remoteAddr)
+
+	suite.conn = newConnProxyProtocol(suite.sourceConnMock, suite.targetConnMock)
+}
+
+func (suite *ConnProxyProtocolTestSuite) TestRead() {
+	value := []byte{1, 2, 3, 4, 5}
+	toRead := make([]byte, len(value))
+
+	suite.targetConnMock.
+		On("Read", mock.AnythingOfType("[]uint8")).
+		Once().
+		Return(len(toRead), nil).
+		Run(func(args mock.Arguments) {
+			arr := args.Get(0).([]byte)
+			copy(arr, value)
+		})
+
+	n, err := suite.conn.Read(toRead)
+	suite.Equal(len(value), n)
+	suite.NoError(err)
+	suite.Equal(value, toRead)
+}
+
+func (suite *ConnProxyProtocolTestSuite) TestWrite() {
+	value := []byte{1, 2, 3, 4, 5}
+	buf := &bytes.Buffer{}
+	bufReader := bufio.NewReader(buf)
+
+	suite.targetConnMock.
+		On("Write", mock.AnythingOfType("[]uint8")).
+		Return(28, nil).
+		Run(func(args mock.Arguments) {
+			arr := args.Get(0).([]byte)
+			buf.Write(arr)
+		})
+
+	_, err := suite.conn.Write(value)
+	suite.NoError(err)
+
+	header, err := proxyproto.Read(bufReader)
+	suite.NoError(err)
+
+	sourceAddr, destAddr, ok := header.TCPAddrs()
+	suite.True(ok)
+	suite.Equal(suite.sourceConnMock.RemoteAddr(), sourceAddr)
+	suite.Equal(suite.targetConnMock.RemoteAddr(), destAddr)
+
+	read, _ := io.ReadAll(bufReader)
+	suite.Equal(value, read)
+
+	_, err = suite.conn.Write(value)
+	suite.NoError(err)
+
+	read, _ = io.ReadAll(bufReader)
+	suite.Equal(value, read)
+}
+
+func (suite *ConnProxyProtocolTestSuite) TearDownTest() {
+	suite.sourceConnMock.AssertExpectations(suite.T())
+	suite.targetConnMock.AssertExpectations(suite.T())
+}
+
 func TestConnTraffic(t *testing.T) {
 	t.Parallel()
 	suite.Run(t, &ConnTrafficTestSuite{})
@@ -209,3 +300,8 @@ func TestConnRewind(t *testing.T) {
 	t.Parallel()
 	suite.Run(t, &ConnRewindTestSuite{})
 }
+
+func TestConnProxyProtocol(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &ConnProxyProtocolTestSuite{})
+}

mtglib/proxy.go

@@ -24,12 +24,13 @@ type Proxy struct {
 	ctxCancel       context.CancelFunc
 	streamWaitGroup sync.WaitGroup
 
-	allowFallbackOnUnknownDC bool
-	tolerateTimeSkewness     time.Duration
-	domainFrontingPort       int
-	domainFrontingIP         string
-	workerPool               *ants.PoolWithFunc
-	telegram                 *dc.Telegram
+	allowFallbackOnUnknownDC    bool
+	tolerateTimeSkewness        time.Duration
+	domainFrontingPort          int
+	domainFrontingIP            string
+	domainFrontingProxyProtocol bool
+	workerPool                  *ants.PoolWithFunc
+	telegram                    *dc.Telegram
 	configUpdater            *dc.PublicConfigUpdater
 	clientObfuscatror        obfuscation.Obfuscator
 
@@ -285,6 +286,10 @@ func (p *Proxy) doDomainFronting(ctx *streamContext, conn *connRewind) {
 		return
 	}
 
+	if p.domainFrontingProxyProtocol {
+		frontConn = newConnProxyProtocol(ctx.clientConn, frontConn)
+	}
+
 	frontConn = connTraffic{
 		Conn:     frontConn,
 		ctx:      ctx,
@@ -316,20 +321,20 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 	updatersLogger := logger.Named("telegram-updaters")
 
 	proxy := &Proxy{
-		ctx:                      ctx,
-		ctxCancel:                cancel,
-		secret:                   opts.Secret,
-		network:                  opts.Network,
-		antiReplayCache:          opts.AntiReplayCache,
-		blocklist:                opts.IPBlocklist,
-		allowlist:                opts.IPAllowlist,
-		eventStream:              opts.EventStream,
+		ctx:                         ctx,
+		ctxCancel:                   cancel,
+		secret:                      opts.Secret,
+		network:                     opts.Network,
+		antiReplayCache:             opts.AntiReplayCache,
+		blocklist:                   opts.IPBlocklist,
+		allowlist:                   opts.IPAllowlist,
+		eventStream:                 opts.EventStream,
 		logger:                   logger,
-		domainFrontingPort:       opts.getDomainFrontingPort(),
-		domainFrontingIP:         opts.DomainFrontingIP,
-		tolerateTimeSkewness:     opts.getTolerateTimeSkewness(),
-		allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC,
-		telegram:                 tg,
+		domainFrontingPort:          opts.getDomainFrontingPort(),
+		domainFrontingIP:            opts.DomainFrontingIP,
+		tolerateTimeSkewness:        opts.getTolerateTimeSkewness(),
+		allowFallbackOnUnknownDC:    opts.AllowFallbackOnUnknownDC,
+		telegram:                    tg,
 		configUpdater: dc.NewPublicConfigUpdater(
 			tg,
 			updatersLogger.Named("public-config"),
@@ -338,6 +343,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		clientObfuscatror: obfuscation.Obfuscator{
 			Secret: opts.Secret.Key[:],
 		},
+		domainFrontingProxyProtocol: opts.DomainFrontingProxyProtocol,
 	}
 
 	proxy.configUpdater.Run(ctx, dc.PublicConfigUpdateURLv4, "tcp4")

mtglib/proxy_opts.go

@@ -102,6 +102,14 @@ type ProxyOpts struct {
 	// This is an optional setting.
 	DomainFrontingIP string
 
+	// DomainFrontingProxyProtocol is used if communication between upstream
+	// endpoint and mtg supports proxy protocol. This is useful in case
+	// if mtg is also placed behind load balancer, and this will make
+	// fronting webserver to know about real IP addresses
+	//
+	// This is an optional setting.
+	DomainFrontingProxyProtocol bool
+
 	// AllowFallbackOnUnknownDC defines how proxy behaves if unknown DC was
 	// requested. If this setting is set to false, then such connection will be
 	// rejected. Otherwise, proxy will chose any DC.