Merge pull request #416 from dolonet/fix/domain-fronting-idle-timeout

fix: apply idle timeout to domain fronting relay

internal/cli/run_proxy.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"time"
 
 	"github.com/9seconds/mtg/v2/antireplay"
 	"github.com/9seconds/mtg/v2/events"
@@ -262,6 +263,7 @@ func runProxy(conf *config.Config, version string) error { //nolint: funlen
 
 		AllowFallbackOnUnknownDC: conf.AllowFallbackOnUnknownDC.Get(false),
 		TolerateTimeSkewness:     conf.TolerateTimeSkewness.Value,
+		IdleTimeout:              conf.Network.Timeout.Idle.Get(time.Minute),
 
 		DoppelGangerURLs:    doppelGangerURLs,
 		DoppelGangerPerRaid: conf.Defense.Doppelganger.Repeats.Get(mtglib.DoppelGangerPerRaid),

mtglib/conns.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"time"
 
 	"github.com/9seconds/mtg/v2/essentials"
 	"github.com/pires/go-proxyproto"
@@ -95,3 +96,21 @@ func newConnProxyProtocol(source, target essentials.Conn) *connProxyProtocol {
 		sourceAddr: source.RemoteAddr(),
 	}
 }
+
+type connIdleTimeout struct {
+	essentials.Conn
+
+	timeout time.Duration
+}
+
+func (c connIdleTimeout) Read(b []byte) (int, error) {
+	c.SetReadDeadline(time.Now().Add(c.timeout)) //nolint: errcheck
+
+	return c.Conn.Read(b) //nolint: wrapcheck
+}
+
+func (c connIdleTimeout) Write(b []byte) (int, error) {
+	c.SetWriteDeadline(time.Now().Add(c.timeout)) //nolint: errcheck
+
+	return c.Conn.Write(b) //nolint: wrapcheck
+}

mtglib/proxy.go

@@ -27,6 +27,7 @@ type Proxy struct {
 
 	allowFallbackOnUnknownDC    bool
 	tolerateTimeSkewness        time.Duration
+	idleTimeout                 time.Duration
 	domainFrontingPort          int
 	domainFrontingIP            string
 	domainFrontingProxyProtocol bool
@@ -151,6 +152,7 @@ func (p *Proxy) Serve(listener net.Listener) error {
 		case errors.Is(err, ants.ErrPoolClosed):
 			return nil
 		case errors.Is(err, ants.ErrPoolOverload):
+			conn.Close() //nolint: errcheck
 			logger.Info("connection was concurrency limited")
 			p.eventStream.Send(p.ctx, NewEventConcurrencyLimited())
 		}
@@ -306,8 +308,8 @@ func (p *Proxy) doDomainFronting(ctx *streamContext, conn *connRewind) {
 	relay.Relay(
 		ctx,
 		ctx.logger.Named("domain-fronting"),
-		frontConn,
-		conn,
+		connIdleTimeout{Conn: frontConn, timeout: p.idleTimeout},
+		connIdleTimeout{Conn: conn, timeout: p.idleTimeout},
 	)
 }
 
@@ -339,6 +341,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		domainFrontingPort:       opts.getDomainFrontingPort(),
 		domainFrontingIP:         opts.DomainFrontingIP,
 		tolerateTimeSkewness:     opts.getTolerateTimeSkewness(),
+		idleTimeout:              opts.getIdleTimeout(),
 		allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC,
 		telegram:                 tg,
 		doppelGanger: doppel.NewGanger(

mtglib/proxy_opts.go

@@ -216,6 +216,14 @@ func (p ProxyOpts) getPreferIP() string {
 	return p.PreferIP
 }
 
+func (p ProxyOpts) getIdleTimeout() time.Duration {
+	if p.IdleTimeout == 0 {
+		return time.Minute
+	}
+
+	return p.IdleTimeout
+}
+
 func (p ProxyOpts) getLogger(name string) Logger {
 	return p.Logger.Named(name)
 }