diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go index 643f645..c57374b 100644 --- a/internal/cli/run_proxy.go +++ b/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), diff --git a/mtglib/conns.go b/mtglib/conns.go index 8219b75..12f19d1 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.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 +} diff --git a/mtglib/proxy.go b/mtglib/proxy.go index 44d426f..fe5eb8d 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 @@ -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( diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index 3d07fbd..102b279 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 time.Minute + } + + return p.IdleTimeout +} + func (p ProxyOpts) getLogger(name string) Logger { return p.Logger.Named(name) }