diff --git a/dialer/tcp/dialer.go b/dialer/tcp/dialer.go index 83af2539..79bbbf99 100644 --- a/dialer/tcp/dialer.go +++ b/dialer/tcp/dialer.go @@ -47,6 +47,18 @@ func (d *tcpDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOp conn, err := options.Dialer.Dial(ctx, "tcp", addr) if err != nil { d.logger.Error(err) + return conn, err + } + + if d.md.keepalive { + if tc, ok := conn.(*net.TCPConn); ok { + tc.SetKeepAliveConfig(net.KeepAliveConfig{ + Enable: true, + Idle: d.md.keepaliveIdle, + Interval: d.md.keepaliveInterval, + Count: d.md.keepaliveCount, + }) + } } conn = proxyproto.WrapClientConn( diff --git a/dialer/tcp/metadata.go b/dialer/tcp/metadata.go index 34e37894..03810a74 100644 --- a/dialer/tcp/metadata.go +++ b/dialer/tcp/metadata.go @@ -4,6 +4,7 @@ import ( "time" md "github.com/go-gost/core/metadata" + mdutil "github.com/go-gost/x/metadata/util" ) const ( @@ -16,8 +17,17 @@ const ( type metadata struct { dialTimeout time.Duration + + keepalive bool + keepaliveIdle time.Duration + keepaliveInterval time.Duration + keepaliveCount int } func (d *tcpDialer) parseMetadata(md md.Metadata) (err error) { + d.md.keepalive = mdutil.GetBool(md, "keepalive") + d.md.keepaliveIdle = mdutil.GetDuration(md, "keepalive.idle") + d.md.keepaliveInterval = mdutil.GetDuration(md, "keepalive.interval") + d.md.keepaliveCount = mdutil.GetInt(md, "keepalive.count") return } diff --git a/listener/tcp/listener.go b/listener/tcp/listener.go index 9d2012b8..c63e3fde 100644 --- a/listener/tcp/listener.go +++ b/listener/tcp/listener.go @@ -61,6 +61,17 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { return } + if l.md.keepalive { + ln = newKeepaliveListener(ln, net.KeepAliveConfig{ + Enable: true, + Idle: l.md.keepaliveIdle, + Interval: l.md.keepaliveInterval, + Count: l.md.keepaliveCount, + }) + l.logger.Debugf("tcp keepalive enabled: idle=%v interval=%v count=%d", + l.md.keepaliveIdle, l.md.keepaliveInterval, l.md.keepaliveCount) + } + l.logger.Debugf("pp: %d", l.options.ProxyProtocol) ln = proxyproto.WrapListener(l.options.ProxyProtocol, ln, 10*time.Second) @@ -100,3 +111,25 @@ func (l *tcpListener) Addr() net.Addr { func (l *tcpListener) Close() error { return l.ln.Close() } + +// keepaliveListener applies TCP keepalive settings to every accepted connection +// before the connection enters the middleware chain. +type keepaliveListener struct { + net.Listener + cfg net.KeepAliveConfig +} + +func newKeepaliveListener(ln net.Listener, cfg net.KeepAliveConfig) net.Listener { + return &keepaliveListener{Listener: ln, cfg: cfg} +} + +func (l *keepaliveListener) Accept() (net.Conn, error) { + conn, err := l.Listener.Accept() + if err != nil { + return nil, err + } + if tc, ok := conn.(*net.TCPConn); ok { + tc.SetKeepAliveConfig(l.cfg) + } + return conn, nil +} diff --git a/listener/tcp/metadata.go b/listener/tcp/metadata.go index 28c91886..cf45f327 100644 --- a/listener/tcp/metadata.go +++ b/listener/tcp/metadata.go @@ -1,16 +1,28 @@ package tcp import ( + "time" + md "github.com/go-gost/core/metadata" mdutil "github.com/go-gost/x/metadata/util" ) type metadata struct { mptcp bool + + keepalive bool + keepaliveIdle time.Duration + keepaliveInterval time.Duration + keepaliveCount int } func (l *tcpListener) parseMetadata(md md.Metadata) (err error) { l.md.mptcp = mdutil.GetBool(md, "mptcp") + l.md.keepalive = mdutil.GetBool(md, "keepalive") + l.md.keepaliveIdle = mdutil.GetDuration(md, "keepalive.idle") + l.md.keepaliveInterval = mdutil.GetDuration(md, "keepalive.interval") + l.md.keepaliveCount = mdutil.GetInt(md, "keepalive.count") + return }