Make TCP keep-alive timeout options configurable

core/option/option.go

@@ -2,6 +2,7 @@ package option
 
 import (
 	"fmt"
+	"time"
 
 	"golang.org/x/time/rate"
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -59,10 +60,30 @@ const (
 	// tcpDefaultReceiveBufferSize is the default size of the receive buffer
 	// for a transport endpoint.
 	tcpDefaultReceiveBufferSize = tcp.DefaultReceiveBufferSize
+
+	// tcpKeepAliveEnabled is the value used to enable or disable keepalives on
+	// the socket.
+	tcpKeepAliveEnabled = true
+
+	// tcpDefaultKeepaliveCount is the maximum number of TCP keep-alive probes to
+	// send before giving up and killing the connection if no response is
+	// obtained from the other end.
+	tcpDefaultKeepaliveCount = 9
+
+	// tcpDefaultKeepaliveIdle specifies the time a connection must remain idle
+	// before the first TCP keepalive packet is sent.
+	tcpDefaultKeepaliveIdle = 60 * time.Second
+
+	// tcpDefaultKeepaliveInterval specifies the interval time between sending
+	// TCP keepalive packets.
+	tcpDefaultKeepaliveInterval = 30 * time.Second
 )
 
 type Option func(*stack.Stack) error
 
+// TCPSocketOption is a function that applies TCP socket-level options to an endpoint.
+type TCPSocketOption func(tcpip.Endpoint) tcpip.Error
+
 // WithDefault sets all default values for stack.
 func WithDefault() Option {
 	return func(s *stack.Stack) error {
@@ -256,3 +277,54 @@ func WithTCPRecovery(v tcpip.TCPRecovery) Option {
 		return nil
 	}
 }
+
+// WithKeepAliveEnabled sets the keep-alive enabled setting
+func WithKeepAliveEnabled(enabled bool) TCPSocketOption {
+	return func(ep tcpip.Endpoint) tcpip.Error {
+		ep.SocketOptions().SetKeepAlive(enabled)
+		return nil
+	}
+}
+
+// WithTCPKeepaliveIdleTime sets the TCP keepalive idle time.
+func WithTCPKeepaliveIdleTime(idle time.Duration) TCPSocketOption {
+	return func(ep tcpip.Endpoint) tcpip.Error {
+		opt := tcpip.KeepaliveIdleOption(idle)
+		return ep.SetSockOpt(&opt)
+	}
+}
+
+// WithTCPKeepaliveInterval sets the TCP keepalive interval.
+func WithTCPKeepaliveInterval(interval time.Duration) TCPSocketOption {
+	return func(ep tcpip.Endpoint) tcpip.Error {
+		opt := tcpip.KeepaliveIntervalOption(interval)
+		return ep.SetSockOpt(&opt)
+	}
+}
+
+// WithTCPKeepaliveCount sets the TCP keepalive count.
+func WithTCPKeepaliveCount(count int) TCPSocketOption {
+	return func(ep tcpip.Endpoint) tcpip.Error {
+		return ep.SetSockOptInt(tcpip.KeepaliveCountOption, count)
+	}
+}
+
+// WithDefaultTCPKeepaliveEnabled sets the default TCP keepalive enabled setting.
+func WithDefaultTCPKeepaliveEnabled() TCPSocketOption {
+	return WithKeepAliveEnabled(tcpKeepAliveEnabled)
+}
+
+// WithDefaultTCPKeepaliveIdleTime sets the default TCP keepalive idle time.
+func WithDefaultTCPKeepaliveIdleTime() TCPSocketOption {
+	return WithTCPKeepaliveIdleTime(tcpDefaultKeepaliveIdle)
+}
+
+// WithDefaultTCPKeepaliveInterval sets the default TCP keepalive interval.
+func WithDefaultTCPKeepaliveInterval() TCPSocketOption {
+	return WithTCPKeepaliveInterval(tcpDefaultKeepaliveInterval)
+}
+
+// WithDefaultTCPKeepaliveCount sets the default TCP keepalive count.
+func WithDefaultTCPKeepaliveCount() TCPSocketOption {
+	return WithTCPKeepaliveCount(tcpDefaultKeepaliveCount)
+}

core/stack.go

@@ -31,6 +31,9 @@ type Config struct {
 	// Options are supplement options to apply settings
 	// for the internal stack.
 	Options []option.Option
+
+	// TCPSocketOptions are TCP socket-level options to apply to each TCP endpoint.
+	TCPSocketOptions []option.TCPSocketOption
 }
 
 // CreateStack creates *stack.Stack with given config.
@@ -61,7 +64,7 @@ func CreateStack(cfg *Config) (*stack.Stack, error) {
 		// before creating NIC, otherwise NIC would dispatch packets
 		// to stack and cause race condition.
 		// Initiate transport protocol (TCP/UDP) with given handler.
-		withTCPHandler(cfg.TransportHandler.HandleTCP),
+		withTCPHandler(cfg.TransportHandler.HandleTCP, cfg.TCPSocketOptions),
 		withUDPHandler(cfg.TransportHandler.HandleUDP),
 
 		// Create stack NIC and then bind link endpoint to it.

core/tcp.go

@@ -1,8 +1,6 @@
 package core
 
 import (
-	"time"
-
 	glog "gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
@@ -23,25 +21,9 @@ const (
 	// maxConnAttempts specifies the maximum number
 	// of in-flight tcp connection attempts.
 	maxConnAttempts = 2 << 10
-
-	// tcpKeepaliveCount is the maximum number of
-	// TCP keep-alive probes to send before giving up
-	// and killing the connection if no response is
-	// obtained from the other end.
-	tcpKeepaliveCount = 9
-
-	// tcpKeepaliveIdle specifies the time a connection
-	// must remain idle before the first TCP keepalive
-	// packet is sent. Once this time is reached,
-	// tcpKeepaliveInterval option is used instead.
-	tcpKeepaliveIdle = 60 * time.Second
-
-	// tcpKeepaliveInterval specifies the interval
-	// time between sending TCP keepalive packets.
-	tcpKeepaliveInterval = 30 * time.Second
 )
 
-func withTCPHandler(handle func(adapter.TCPConn)) option.Option {
+func withTCPHandler(handle func(adapter.TCPConn), tcpSockOpts []option.TCPSocketOption) option.Option {
 	return func(s *stack.Stack) error {
 		tcpForwarder := tcp.NewForwarder(s, defaultWndSize, maxConnAttempts, func(r *tcp.ForwarderRequest) {
 			var (
@@ -67,7 +49,7 @@ func withTCPHandler(handle func(adapter.TCPConn)) option.Option {
 			}
 			defer r.Complete(false)
 
-			err = setSocketOptions(s, ep)
+			err = setSocketOptions(s, ep, tcpSockOpts)
 
 			conn := &tcpConn{
 				TCPConn: gonet.NewTCPConn(&wq, ep),
@@ -80,22 +62,12 @@ func withTCPHandler(handle func(adapter.TCPConn)) option.Option {
 	}
 }
 
-func setSocketOptions(s *stack.Stack, ep tcpip.Endpoint) tcpip.Error {
-	{ /* TCP keepalive options */
-		ep.SocketOptions().SetKeepAlive(true)
-
-		idle := tcpip.KeepaliveIdleOption(tcpKeepaliveIdle)
-		if err := ep.SetSockOpt(&idle); err != nil {
-			return err
-		}
-
-		interval := tcpip.KeepaliveIntervalOption(tcpKeepaliveInterval)
-		if err := ep.SetSockOpt(&interval); err != nil {
-			return err
-		}
-
-		if err := ep.SetSockOptInt(tcpip.KeepaliveCountOption, tcpKeepaliveCount); err != nil {
-			return err
+func setSocketOptions(s *stack.Stack, ep tcpip.Endpoint, tcpSockOpts []option.TCPSocketOption) tcpip.Error {
+	{ /* Apply socket options (defaults + custom overrides)*/
+		for _, opt := range tcpSockOpts {
+			if err := opt(ep); err != nil {
+				return err
+			}
 		}
 	}
 	{ /* TCP recv/send buffer size */

engine/engine.go

@@ -205,6 +205,9 @@ func netstack(k *Key) (err error) {
 	}
 
 	var opts []option.Option
+	tcpSocketOpts := []option.TCPSocketOption{
+		option.WithDefaultTCPKeepaliveEnabled(),
+	}
 	if k.TCPModerateReceiveBuffer {
 		opts = append(opts, option.WithTCPModerateReceiveBuffer(true))
 	}
@@ -225,11 +228,33 @@ func netstack(k *Key) (err error) {
 		opts = append(opts, option.WithTCPReceiveBufferSize(int(size)))
 	}
 
+	if k.TCPKeepaliveIdleTime > 0 {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithTCPKeepaliveIdleTime(k.TCPKeepaliveIdleTime))
+		log.Infof("[STACK] keepalive idle time: %v", k.TCPKeepaliveIdleTime)
+	} else {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithDefaultTCPKeepaliveIdleTime())
+	}
+
+	if k.TCPKeepaliveInterval > 0 {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithTCPKeepaliveInterval(k.TCPKeepaliveInterval))
+		log.Infof("[STACK] keepalive interval: %v", k.TCPKeepaliveInterval)
+	} else {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithDefaultTCPKeepaliveInterval())
+	}
+
+	if k.TCPKeepaliveCount > 0 {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithTCPKeepaliveCount(k.TCPKeepaliveCount))
+		log.Infof("[STACK] keepalive count: %d", k.TCPKeepaliveCount)
+	} else {
+		tcpSocketOpts = append(tcpSocketOpts, option.WithDefaultTCPKeepaliveCount())
+	}
+
 	if _defaultStack, err = core.CreateStack(&core.Config{
 		LinkEndpoint:     _defaultDevice,
 		TransportHandler: tunnel.T(),
 		MulticastGroups:  multicastGroups,
 		Options:          opts,
+		TCPSocketOptions: tcpSocketOpts,
 	}); err != nil {
 		return err
 	}

engine/key.go

@@ -10,6 +10,9 @@ type Key struct {
 	Device                   string        `yaml:"device"`
 	LogLevel                 string        `yaml:"loglevel"`
 	Interface                string        `yaml:"interface"`
+	TCPKeepaliveCount        int           `yaml:"tcp-keepalive-count"`
+	TCPKeepaliveIdleTime     time.Duration `yaml:"tcp-keepalive-idle-time"`
+	TCPKeepaliveInterval     time.Duration `yaml:"tcp-keepalive-interval"`
 	TCPModerateReceiveBuffer bool          `yaml:"tcp-moderate-receive-buffer"`
 	TCPSendBufferSize        string        `yaml:"tcp-send-buffer-size"`
 	TCPReceiveBufferSize     string        `yaml:"tcp-receive-buffer-size"`

main.go

@@ -38,6 +38,11 @@ func init() {
 	flag.StringVar(&key.MulticastGroups, "multicast-groups", "", "Set multicast groups, separated by commas")
 	flag.StringVar(&key.TUNPreUp, "tun-pre-up", "", "Execute a command before TUN device setup")
 	flag.StringVar(&key.TUNPostUp, "tun-post-up", "", "Execute a command after TUN device setup")
+
+	flag.DurationVar(&key.TCPKeepaliveIdleTime, "tcp-keepalive-idle-time", 0, "TCP keepalive idle time before first probe (e.g., 60s, 2m)")
+	flag.DurationVar(&key.TCPKeepaliveInterval, "tcp-keepalive-interval", 0, "TCP keepalive probe interval (e.g., 30s, 1m)")
+	flag.IntVar(&key.TCPKeepaliveCount, "tcp-keepalive-count", 0, "TCP keepalive probe count before giving up")
+
 	flag.BoolVar(&versionFlag, "version", false, "Show version and then quit")
 	flag.Parse()
 }