Refactor: core.CreateStackWithOptions

core/tcp.go

@@ -0,0 +1,94 @@
+package core
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/xjasonlyu/tun2socks/v2/core/adapter"
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
+
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
+	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
+	"gvisor.dev/gvisor/pkg/waiter"
+)
+
+const (
+	// defaultWndSize if set to zero, the default
+	// receive window buffer size is used instead.
+	defaultWndSize = 0
+
+	// 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 adapter.TCPHandleFunc) option.Option {
+	return func(s *stack.Stack) error {
+		tcpForwarder := tcp.NewForwarder(s, defaultWndSize, maxConnAttempts, func(r *tcp.ForwarderRequest) {
+			var wq waiter.Queue
+			ep, err := r.CreateEndpoint(&wq)
+			if err != nil {
+				// RST: prevent potential half-open TCP connection leak.
+				r.Complete(true)
+				return
+			}
+			defer r.Complete(false)
+
+			setKeepalive(ep)
+
+			conn := &tcpConn{
+				TCPConn: gonet.NewTCPConn(&wq, ep),
+				id:      r.ID(),
+			}
+			handle(conn)
+		})
+		s.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
+		return nil
+	}
+}
+
+func setKeepalive(ep tcpip.Endpoint) error {
+	ep.SocketOptions().SetKeepAlive(true)
+
+	idle := tcpip.KeepaliveIdleOption(tcpKeepaliveIdle)
+	if err := ep.SetSockOpt(&idle); err != nil {
+		return fmt.Errorf("set keepalive idle: %s", err)
+	}
+
+	interval := tcpip.KeepaliveIntervalOption(tcpKeepaliveInterval)
+	if err := ep.SetSockOpt(&interval); err != nil {
+		return fmt.Errorf("set keepalive interval: %s", err)
+	}
+
+	if err := ep.SetSockOptInt(tcpip.KeepaliveCountOption, tcpKeepaliveCount); err != nil {
+		return fmt.Errorf("set keepalive count: %s", err)
+	}
+	return nil
+}
+
+type tcpConn struct {
+	*gonet.TCPConn
+	id stack.TransportEndpointID
+}
+
+func (c *tcpConn) ID() *stack.TransportEndpointID {
+	return &c.id
+}