diff --git a/core/option/option.go b/core/option/option.go index 2076fd58..6eda0c24 100644 --- a/core/option/option.go +++ b/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) +} diff --git a/core/stack.go b/core/stack.go index 1f388832..236c7502 100644 --- a/core/stack.go +++ b/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. diff --git a/core/tcp.go b/core/tcp.go index 5c5ee0f6..f207d24c 100644 --- a/core/tcp.go +++ b/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 */ diff --git a/engine/engine.go b/engine/engine.go index 17ce5554..d70b903a 100644 --- a/engine/engine.go +++ b/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 } diff --git a/engine/key.go b/engine/key.go index 5a86d53a..8f94106e 100644 --- a/engine/key.go +++ b/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"` diff --git a/main.go b/main.go index ac425905..2816ec31 100644 --- a/main.go +++ b/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() }