Add support for domain fronting proxy protocol

example.config.toml

@@ -67,6 +67,10 @@ domain-fronting-port = 443
 # default value is not set (DNS resolution is used).
 # domain-fronting-ip = "142.250.185.112"
 
+# This makes a communication between both fronting website and mtg to use
+# proxy protocol.
+domain-fronting-proxy-protocol = false
+
 # FakeTLS can compare timestamps to prevent probes. Each message has
 # encrypted timestamp. So, mtg can compare this timestamp and decide if
 # we need to proceed with connection or not.

internal/cli/run_proxy.go

@@ -250,10 +250,11 @@ func runProxy(conf *config.Config, version string) error { //nolint: funlen
 		IPAllowlist:     allowlist,
 		EventStream:     eventStream,
 
-		Secret:             conf.Secret,
-		DomainFrontingPort: conf.DomainFrontingPort.Get(mtglib.DefaultDomainFrontingPort),
-		DomainFrontingIP:   conf.DomainFrontingIP.String(),
-		PreferIP:           conf.PreferIP.Get(mtglib.DefaultPreferIP),
+		Secret:                      conf.Secret,
+		DomainFrontingPort:          conf.DomainFrontingPort.Get(mtglib.DefaultDomainFrontingPort),
+		DomainFrontingIP:            conf.DomainFrontingIP.String(),
+		DomainFrontingProxyProtocol: conf.DomainFrontingProxyProtocol.Get(false),
+		PreferIP:                    conf.PreferIP.Get(mtglib.DefaultPreferIP),
 
 		AllowFallbackOnUnknownDC: conf.AllowFallbackOnUnknownDC.Get(false),
 		TolerateTimeSkewness:     conf.TolerateTimeSkewness.Value,

internal/config/config.go

@@ -21,17 +21,18 @@ type ListConfig struct {
 }
 
 type Config struct {
-	Debug                    TypeBool        `json:"debug"`
-	AllowFallbackOnUnknownDC TypeBool        `json:"allowFallbackOnUnknownDc"`
-	Secret                   mtglib.Secret   `json:"secret"`
-	BindTo                   TypeHostPort    `json:"bindTo"`
-	ProxyProtocolListener    TypeBool        `json:"proxyProtocolListener"`
-	PreferIP                 TypePreferIP    `json:"preferIp"`
-	DomainFrontingPort       TypePort        `json:"domainFrontingPort"`
-	DomainFrontingIP         TypeIP          `json:"domainFrontingIp"`
-	TolerateTimeSkewness     TypeDuration    `json:"tolerateTimeSkewness"`
-	Concurrency              TypeConcurrency `json:"concurrency"`
-	Defense                  struct {
+	Debug                       TypeBool        `json:"debug"`
+	AllowFallbackOnUnknownDC    TypeBool        `json:"allowFallbackOnUnknownDc"`
+	Secret                      mtglib.Secret   `json:"secret"`
+	BindTo                      TypeHostPort    `json:"bindTo"`
+	ProxyProtocolListener       TypeBool        `json:"proxyProtocolListener"`
+	PreferIP                    TypePreferIP    `json:"preferIp"`
+	DomainFrontingPort          TypePort        `json:"domainFrontingPort"`
+	DomainFrontingIP            TypeIP          `json:"domainFrontingIp"`
+	DomainFrontingProxyProtocol TypeBool        `json:"domainFrontingProxyProtocol"`
+	TolerateTimeSkewness        TypeDuration    `json:"tolerateTimeSkewness"`
+	Concurrency                 TypeConcurrency `json:"concurrency"`
+	Defense                     struct {
 		AntiReplay struct {
 			Optional
 

internal/config/parse.go

@@ -9,17 +9,18 @@ import (
 )
 
 type tomlConfig struct {
-	Debug                    bool   `toml:"debug" json:"debug,omitempty"`
-	AllowFallbackOnUnknownDC bool   `toml:"allow-fallback-on-unknown-dc" json:"allowFallbackOnUnknownDc,omitempty"`
-	Secret                   string `toml:"secret" json:"secret"`
-	BindTo                   string `toml:"bind-to" json:"bindTo"`
-	ProxyProtocolListener    bool   `toml:"proxy-protocol-listener" json:"proxyProtocolListener"`
-	PreferIP                 string `toml:"prefer-ip" json:"preferIp,omitempty"`
-	DomainFrontingPort       uint   `toml:"domain-fronting-port" json:"domainFrontingPort,omitempty"`
-	DomainFrontingIP         string `toml:"domain-fronting-ip" json:"domainFrontingIp,omitempty"`
-	TolerateTimeSkewness     string `toml:"tolerate-time-skewness" json:"tolerateTimeSkewness,omitempty"`
-	Concurrency              uint   `toml:"concurrency" json:"concurrency,omitempty"`
-	Defense                  struct {
+	Debug                       bool   `toml:"debug" json:"debug,omitempty"`
+	AllowFallbackOnUnknownDC    bool   `toml:"allow-fallback-on-unknown-dc" json:"allowFallbackOnUnknownDc,omitempty"`
+	Secret                      string `toml:"secret" json:"secret"`
+	BindTo                      string `toml:"bind-to" json:"bindTo"`
+	ProxyProtocolListener       bool   `toml:"proxy-protocol-listener" json:"proxyProtocolListener"`
+	PreferIP                    string `toml:"prefer-ip" json:"preferIp,omitempty"`
+	DomainFrontingPort          uint   `toml:"domain-fronting-port" json:"domainFrontingPort,omitempty"`
+	DomainFrontingIP            string `toml:"domain-fronting-ip" json:"domainFrontingIp,omitempty"`
+	DomainFrontingProxyProtocol bool   `toml:"domain-fronting-proxy-protocol" json:"domainFrontingProxyProtocol,omitempty"`
+	TolerateTimeSkewness        string `toml:"tolerate-time-skewness" json:"tolerateTimeSkewness,omitempty"`
+	Concurrency                 uint   `toml:"concurrency" json:"concurrency,omitempty"`
+	Defense                     struct {
 		AntiReplay struct {
 			Enabled   bool    `toml:"enabled" json:"enabled,omitempty"`
 			MaxSize   string  `toml:"max-size" json:"maxSize,omitempty"`

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.