Refactor: core.CreateStackWithOptions

core/adapter/handler.go

@@ -3,6 +3,12 @@ package adapter
 // Handler is a TCP/UDP connection handler that implements
 // HandleTCPConn and HandleUDPConn methods.
 type Handler interface {
-	HandleTCPConn(TCPConn)
+	HandleTCP(TCPConn)
-	HandleUDPConn(UDPConn)
+	HandleUDP(UDPConn)
 }
+
+// TCPHandleFunc handles incoming TCP connection.
+type TCPHandleFunc func(TCPConn)
+
+// UDPHandleFunc handles incoming UDP connection.
+type UDPHandleFunc func(UDPConn)

core/nic.go

@@ -1,16 +1,15 @@
-package stack
+package core
 
 import (
 	"fmt"
 
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
+
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 const (
-	// defaultNICID is the ID of default NIC used by DefaultStack.
-	defaultNICID tcpip.NICID = 0x01
-
 	// nicPromiscuousModeEnabled is the value used by stack to enable
 	// or disable NIC's promiscuous mode.
 	nicPromiscuousModeEnabled = true
@@ -21,9 +20,9 @@ const (
 )
 
 // withCreatingNIC creates NIC for stack.
-func withCreatingNIC(ep stack.LinkEndpoint) Option {
+func withCreatingNIC(nicID tcpip.NICID, ep stack.LinkEndpoint) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		if err := s.CreateNICWithOptions(s.nicID, ep,
+		if err := s.CreateNICWithOptions(nicID, ep,
 			stack.NICOptions{
 				Disabled: false,
 				// If no queueing discipline was specified
@@ -37,21 +36,21 @@ func withCreatingNIC(ep stack.LinkEndpoint) Option {
 	}
 }
 
-// withPromiscuousMode sets promiscuous mode in the given NIC.
+// withPromiscuousMode sets promiscuous mode in the given NICs.
-func withPromiscuousMode(v bool) Option {
+func withPromiscuousMode(nicID tcpip.NICID, v bool) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		if err := s.SetPromiscuousMode(s.nicID, v); err != nil {
+		if err := s.SetPromiscuousMode(nicID, v); err != nil {
 			return fmt.Errorf("set promiscuous mode: %s", err)
 		}
 		return nil
 	}
 }
 
-// withSpoofing sets address spoofing in the given NIC, allowing
+// withSpoofing sets address spoofing in the given NICs, allowing
 // endpoints to bind to any address in the NIC.
-func withSpoofing(v bool) Option {
+func withSpoofing(nicID tcpip.NICID, v bool) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		if err := s.SetSpoofing(s.nicID, v); err != nil {
+		if err := s.SetSpoofing(nicID, v); err != nil {
 			return fmt.Errorf("set spoofing: %s", err)
 		}
 		return nil

core/option/option.go

@@ -1,4 +1,4 @@
-package stack
+package option
 
 import (
 	"fmt"
@@ -7,6 +7,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
 )
 
@@ -56,11 +57,11 @@ const (
 	tcpDefaultBufferSize = 212 << 10 // 212 KiB
 )
 
-type Option func(*Stack) error
+type Option func(*stack.Stack) error
 
 // WithDefault sets all default values for stack.
 func WithDefault() Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opts := []Option{
 			WithDefaultTTL(defaultTimeToLive),
 			WithForwarding(ipForwardingEnabled),
@@ -110,7 +111,7 @@ func WithDefault() Option {
 
 // WithDefaultTTL sets the default TTL used by stack.
 func WithDefaultTTL(ttl uint8) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opt := tcpip.DefaultTTLOption(ttl)
 		if err := s.SetNetworkProtocolOption(ipv4.ProtocolNumber, &opt); err != nil {
 			return fmt.Errorf("set ipv4 default TTL: %s", err)
@@ -124,7 +125,7 @@ func WithDefaultTTL(ttl uint8) Option {
 
 // WithForwarding sets packet forwarding between NICs for IPv4 & IPv6.
 func WithForwarding(v bool) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		if err := s.SetForwardingDefaultAndAllNICs(ipv4.ProtocolNumber, v); err != nil {
 			return fmt.Errorf("set ipv4 forwarding: %s", err)
 		}
@@ -138,7 +139,7 @@ func WithForwarding(v bool) Option {
 // WithICMPBurst sets the number of ICMP messages that can be sent
 // in a single burst.
 func WithICMPBurst(burst int) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		s.SetICMPBurst(burst)
 		return nil
 	}
@@ -147,7 +148,7 @@ func WithICMPBurst(burst int) Option {
 // WithICMPLimit sets the maximum number of ICMP messages permitted
 // by rate limiter.
 func WithICMPLimit(limit rate.Limit) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		s.SetICMPLimit(limit)
 		return nil
 	}
@@ -155,7 +156,7 @@ func WithICMPLimit(limit rate.Limit) Option {
 
 // WithTCPBufferSizeRange sets the receive and send buffer size range for TCP.
 func WithTCPBufferSizeRange(a, b, c int) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		rcvOpt := tcpip.TCPReceiveBufferSizeRangeOption{Min: a, Default: b, Max: c}
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &rcvOpt); err != nil {
 			return fmt.Errorf("set TCP receive buffer size range: %s", err)
@@ -170,7 +171,7 @@ func WithTCPBufferSizeRange(a, b, c int) Option {
 
 // WithTCPCongestionControl sets the current congestion control algorithm.
 func WithTCPCongestionControl(cc string) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opt := tcpip.CongestionControlOption(cc)
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &opt); err != nil {
 			return fmt.Errorf("set TCP congestion control algorithm: %s", err)
@@ -181,7 +182,7 @@ func WithTCPCongestionControl(cc string) Option {
 
 // WithTCPDelay enables or disables Nagle's algorithm in TCP.
 func WithTCPDelay(v bool) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opt := tcpip.TCPDelayEnabled(v)
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &opt); err != nil {
 			return fmt.Errorf("set TCP delay: %s", err)
@@ -192,7 +193,7 @@ func WithTCPDelay(v bool) Option {
 
 // WithTCPModerateReceiveBuffer sets receive buffer moderation for TCP.
 func WithTCPModerateReceiveBuffer(v bool) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opt := tcpip.TCPModerateReceiveBufferOption(v)
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &opt); err != nil {
 			return fmt.Errorf("set TCP moderate receive buffer: %s", err)
@@ -203,7 +204,7 @@ func WithTCPModerateReceiveBuffer(v bool) Option {
 
 // WithTCPSACKEnabled sets the SACK option for TCP.
 func WithTCPSACKEnabled(v bool) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		opt := tcpip.TCPSACKEnabled(v)
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &opt); err != nil {
 			return fmt.Errorf("set TCP SACK: %s", err)
@@ -214,7 +215,7 @@ func WithTCPSACKEnabled(v bool) Option {
 
 // WithTCPRecovery sets the recovery option for TCP.
 func WithTCPRecovery(v tcpip.TCPRecovery) Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
 		if err := s.SetTransportProtocolOption(tcp.ProtocolNumber, &v); err != nil {
 			return fmt.Errorf("set TCP Recovery: %s", err)
 		}

core/route.go

@@ -1,22 +1,23 @@
-package stack
+package core
 
 import (
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
+
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
-func withICMPHandler() Option {
+func withRouteTable(nicID tcpip.NICID) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		// Add default route table for IPv4 and IPv6.
-		// This will handle all incoming ICMP packets.
 		s.SetRouteTable([]tcpip.Route{
 			{
 				Destination: header.IPv4EmptySubnet,
-				NIC:         s.nicID,
+				NIC:         nicID,
 			},
 			{
 				Destination: header.IPv6EmptySubnet,
-				NIC:         s.nicID,
+				NIC:         nicID,
 			},
 		})
 		return nil

core/stack.go

@@ -0,0 +1,75 @@
+package core
+
+import (
+	"github.com/xjasonlyu/tun2socks/v2/core/adapter"
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
+
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
+	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
+	"gvisor.dev/gvisor/pkg/tcpip/transport/icmp"
+	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
+	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
+)
+
+// CreateStackWithOptions creates *stack.Stack with given options.
+func CreateStackWithOptions(linkEP stack.LinkEndpoint, handler adapter.Handler, opts ...option.Option) (*stack.Stack, error) {
+	s := stack.New(stack.Options{
+		NetworkProtocols: []stack.NetworkProtocolFactory{
+			ipv4.NewProtocol,
+			ipv6.NewProtocol,
+		},
+		TransportProtocols: []stack.TransportProtocolFactory{
+			tcp.NewProtocol,
+			udp.NewProtocol,
+			icmp.NewProtocol4,
+			icmp.NewProtocol6,
+		},
+	})
+
+	// Generate unique NIC id.
+	nicID := tcpip.NICID(s.UniqueID())
+
+	opts = append(opts,
+		// Create stack NIC and then bind link endpoint to it.
+		withCreatingNIC(nicID, linkEP),
+
+		// In the past we did s.AddAddressRange to assign 0.0.0.0/0
+		// onto the interface. We need that to be able to terminate
+		// all the incoming connections - to any ip. AddressRange API
+		// has been removed and the suggested workaround is to use
+		// Promiscuous mode. https://github.com/google/gvisor/issues/3876
+		//
+		// Ref: https://github.com/cloudflare/slirpnetstack/blob/master/stack.go
+		withPromiscuousMode(nicID, nicPromiscuousModeEnabled),
+
+		// Enable spoofing if a stack may send packets from unowned
+		// addresses. This change required changes to some netgophers
+		// since previously, promiscuous mode was enough to let the
+		// netstack respond to all incoming packets regardless of the
+		// packet's destination address. Now that a stack.Route is not
+		// held for each incoming packet, finding a route may fail with
+		// local addresses we don't own but accepted packets for while
+		// in promiscuous mode. Since we also want to be able to send
+		// from any address (in response the received promiscuous mode
+		// packets), we need to enable spoofing.
+		//
+		// Ref: https://github.com/google/gvisor/commit/8c0701462a84ff77e602f1626aec49479c308127
+		withSpoofing(nicID, nicSpoofingEnabled),
+
+		// Add default route table for IPv4 and IPv6. This will handle
+		// all incoming ICMP packets.
+		withRouteTable(nicID),
+
+		// Initiate transport protocol (TCP/UDP) with given handler.
+		withTCPHandler(handler.HandleTCP), withUDPHandler(handler.HandleUDP),
+	)
+
+	for _, opt := range opts {
+		if err := opt(s); err != nil {
+			return nil, err
+		}
+	}
+	return s, nil
+}

core/stack/stack.go

@@ -1,82 +0,0 @@
-// Package stack provides a thin wrapper around a gVisor's stack.
-package stack
-
-import (
-	"github.com/xjasonlyu/tun2socks/v2/core/adapter"
-
-	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
-	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/icmp"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
-	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
-)
-
-type Stack struct {
-	*stack.Stack
-
-	handler adapter.Handler
-	nicID   tcpip.NICID
-}
-
-// New allocates a new *Stack with given options.
-func New(ep stack.LinkEndpoint, handler adapter.Handler, opts ...Option) (*Stack, error) {
-	s := &Stack{
-		Stack: stack.New(stack.Options{
-			NetworkProtocols: []stack.NetworkProtocolFactory{
-				ipv4.NewProtocol,
-				ipv6.NewProtocol,
-			},
-			TransportProtocols: []stack.TransportProtocolFactory{
-				tcp.NewProtocol,
-				udp.NewProtocol,
-				icmp.NewProtocol4,
-				icmp.NewProtocol6,
-			},
-		}),
-
-		handler: handler,
-		nicID:   defaultNICID,
-	}
-
-	opts = append(opts,
-		// Important: We must initiate transport protocol handlers
-		// before creating NIC, otherwise NIC would dispatch packets
-		// to stack and cause race condition.
-		withICMPHandler(), withTCPHandler(), withUDPHandler(),
-
-		// Create stack NIC and then bind link endpoint.
-		withCreatingNIC(ep),
-
-		// In the past we did s.AddAddressRange to assign 0.0.0.0/0
-		// onto the interface. We need that to be able to terminate
-		// all the incoming connections - to any ip. AddressRange API
-		// has been removed and the suggested workaround is to use
-		// Promiscuous mode. https://github.com/google/gvisor/issues/3876
-		//
-		// Ref: https://github.com/cloudflare/slirpnetstack/blob/master/stack.go
-		withPromiscuousMode(nicPromiscuousModeEnabled),
-
-		// Enable spoofing if a stack may send packets from unowned addresses.
-		// This change required changes to some netgophers since previously,
-		// promiscuous mode was enough to let the netstack respond to all
-		// incoming packets regardless of the packet's destination address. Now
-		// that a stack.Route is not held for each incoming packet, finding a route
-		// may fail with local addresses we don't own but accepted packets for
-		// while in promiscuous mode. Since we also want to be able to send from
-		// any address (in response the received promiscuous mode packets), we need
-		// to enable spoofing.
-		//
-		// Ref: https://github.com/google/gvisor/commit/8c0701462a84ff77e602f1626aec49479c308127
-		withSpoofing(nicSpoofingEnabled),
-	)
-
-	for _, opt := range opts {
-		if err := opt(s); err != nil {
-			return nil, err
-		}
-	}
-
-	return s, nil
-}

core/tcp.go

@@ -1,9 +1,12 @@
-package stack
+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"
@@ -37,9 +40,9 @@ const (
 	tcpKeepaliveInterval = 30 * time.Second
 )
 
-func withTCPHandler() Option {
+func withTCPHandler(handle adapter.TCPHandleFunc) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		tcpForwarder := tcp.NewForwarder(s.Stack, defaultWndSize, maxConnAttempts, func(r *tcp.ForwarderRequest) {
+		tcpForwarder := tcp.NewForwarder(s, defaultWndSize, maxConnAttempts, func(r *tcp.ForwarderRequest) {
 			var wq waiter.Queue
 			ep, err := r.CreateEndpoint(&wq)
 			if err != nil {
@@ -55,7 +58,7 @@ func withTCPHandler() Option {
 				TCPConn: gonet.NewTCPConn(&wq, ep),
 				id:      r.ID(),
 			}
-			s.handler.HandleTCPConn(conn)
+			handle(conn)
 		})
 		s.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
 		return nil

core/udp.go

@@ -1,15 +1,18 @@
-package stack
+package core
 
 import (
+	"github.com/xjasonlyu/tun2socks/v2/core/adapter"
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
+
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
 	"gvisor.dev/gvisor/pkg/waiter"
 )
 
-func withUDPHandler() Option {
+func withUDPHandler(handle adapter.UDPHandleFunc) option.Option {
-	return func(s *Stack) error {
+	return func(s *stack.Stack) error {
-		udpForwarder := udp.NewForwarder(s.Stack, func(r *udp.ForwarderRequest) {
+		udpForwarder := udp.NewForwarder(s, func(r *udp.ForwarderRequest) {
 			var wq waiter.Queue
 			ep, err := r.CreateEndpoint(&wq)
 			if err != nil {
@@ -18,10 +21,10 @@ func withUDPHandler() Option {
 			}
 
 			conn := &udpConn{
-				UDPConn: gonet.NewUDPConn(s.Stack, &wq, ep),
+				UDPConn: gonet.NewUDPConn(s, &wq, ep),
 				id:      r.ID(),
 			}
-			s.handler.HandleUDPConn(conn)
+			handle(conn)
 		})
 		s.SetTransportProtocolHandler(udp.ProtocolNumber, udpForwarder.HandlePacket)
 		return nil

engine/engine.go

@@ -5,13 +5,16 @@ import (
 	"net"
 
 	"github.com/xjasonlyu/tun2socks/v2/component/dialer"
+	"github.com/xjasonlyu/tun2socks/v2/core"
 	"github.com/xjasonlyu/tun2socks/v2/core/device"
-	"github.com/xjasonlyu/tun2socks/v2/core/stack"
+	"github.com/xjasonlyu/tun2socks/v2/core/option"
 	_ "github.com/xjasonlyu/tun2socks/v2/dns"
 	"github.com/xjasonlyu/tun2socks/v2/log"
 	"github.com/xjasonlyu/tun2socks/v2/proxy"
 	"github.com/xjasonlyu/tun2socks/v2/stats"
 	"github.com/xjasonlyu/tun2socks/v2/tunnel"
+
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 var _engine = &engine{}
@@ -171,6 +174,6 @@ func (e *engine) applyStack() (err error) {
 		}
 	}()
 
-	e.stack, err = stack.New(e.device, &fakeTunnel{}, stack.WithDefault())
+	e.stack, err = core.CreateStackWithOptions(e.device, &fakeTunnel{}, option.WithDefault())
 	return
 }

engine/tunnel.go

@@ -9,10 +9,10 @@ var _ adapter.Handler = (*fakeTunnel)(nil)
 
 type fakeTunnel struct{}
 
-func (*fakeTunnel) HandleTCPConn(conn adapter.TCPConn) {
+func (*fakeTunnel) HandleTCP(conn adapter.TCPConn) {
 	tunnel.TCPIn() <- conn
 }
 
-func (*fakeTunnel) HandleUDPConn(conn adapter.UDPConn) {
+func (*fakeTunnel) HandleUDP(conn adapter.UDPConn) {
 	tunnel.UDPIn() <- conn
 }