diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go index 643f645..8eabc3d 100644 --- a/internal/cli/run_proxy.go +++ b/internal/cli/run_proxy.go @@ -262,6 +262,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(mtglib.DefaultIdleTimeout), DoppelGangerURLs: doppelGangerURLs, DoppelGangerPerRaid: conf.Defense.Doppelganger.Repeats.Get(mtglib.DoppelGangerPerRaid), diff --git a/mtglib/conns.go b/mtglib/conns.go index 8219b75..e3b60e8 100644 --- a/mtglib/conns.go +++ b/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.Conn.SetReadDeadline(time.Now().Add(c.timeout)) //nolint: errcheck + + return c.Conn.Read(b) //nolint: wrapcheck +} + +func (c connIdleTimeout) Write(b []byte) (int, error) { + c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)) //nolint: errcheck + + return c.Conn.Write(b) //nolint: wrapcheck +} diff --git a/mtglib/proxy.go b/mtglib/proxy.go index be05194..65f52ad 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -27,6 +27,7 @@ type Proxy struct { allowFallbackOnUnknownDC bool tolerateTimeSkewness time.Duration + idleTimeout time.Duration domainFrontingPort int domainFrontingIP string domainFrontingProxyProtocol bool @@ -306,8 +307,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 +340,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( diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index 3d07fbd..beea1c6 100644 --- a/mtglib/proxy_opts.go +++ b/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 DefaultIdleTimeout + } + + return p.IdleTimeout +} + func (p ProxyOpts) getLogger(name string) Logger { return p.Logger.Named(name) }