Add pause support

Author: nekohasekai

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/sagernet/go-tun2socks v1.16.12-0.20220818015926-16cb67876a61
 	github.com/sagernet/gvisor v0.0.0-20230627031050-1ab0276e0dd2
 	github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97
-	github.com/sagernet/sing v0.2.9
+	github.com/sagernet/sing v0.2.10-0.20230831053650-deb9afed5dbf
 	github.com/scjalliance/comshim v0.0.0-20230315213746-5e51f40bd3b9
 	golang.org/x/net v0.14.0
 	golang.org/x/sys v0.11.0

go.sum

@@ -11,8 +11,8 @@ github.com/sagernet/gvisor v0.0.0-20230627031050-1ab0276e0dd2/go.mod h1:1JUiV7nG
 github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97 h1:iL5gZI3uFp0X6EslacyapiRz7LLSJyr4RajF/BhMVyE=
 github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97/go.mod h1:xLnfdiJbSp8rNqYEdIW/6eDO4mVoogml14Bh2hSiFpM=
 github.com/sagernet/sing v0.0.0-20220817130738-ce854cda8522/go.mod h1:QVsS5L/ZA2Q5UhQwLrn0Trw+msNd/NPGEhBKR/ioWiY=
-github.com/sagernet/sing v0.2.9 h1:3wsTz+JG5Wzy65eZnh6AuCrD2QqcRF6Iq6f7ttmJsAo=
-github.com/sagernet/sing v0.2.9/go.mod h1:Ta8nHnDLAwqySzKhGoKk4ZIB+vJ3GTKj7UPrWYvM+4w=
+github.com/sagernet/sing v0.2.10-0.20230831053650-deb9afed5dbf h1:pnmqxDWO/uw4ElokoJiFUVZhCrp80CcENpYagdRL7sI=
+github.com/sagernet/sing v0.2.10-0.20230831053650-deb9afed5dbf/go.mod h1:9uOZwWkhT2Z2WldolLxX34s+1svAX4i4vvz5hy8u1MA=
 github.com/scjalliance/comshim v0.0.0-20230315213746-5e51f40bd3b9 h1:rc/CcqLH3lh8n+csdOuDfP+NuykE0U6AeYSJJHKDgSg=
 github.com/scjalliance/comshim v0.0.0-20230315213746-5e51f40bd3b9/go.mod h1:a/83NAfUXvEuLpmxDssAXxgUgrEy12MId3Wd7OTs76s=
 github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 h1:gga7acRE695APm9hlsSMoOoE65U4/TcqNj90mc69Rlg=

stack_gvisor.go

@@ -23,6 +23,7 @@ import (
 	E "github.com/sagernet/sing/common/exceptions"
 	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
+	"github.com/sagernet/sing/service/pause"
 )
 
 const WithGVisor = true
@@ -37,6 +38,7 @@ type GVisor struct {
 	udpTimeout             int64
 	handler                Handler
 	logger                 logger.Logger
+	pauseManager           pause.Manager
 	stack                  *stack.Stack
 	endpoint               stack.LinkEndpoint
 }
@@ -62,6 +64,7 @@ func NewGVisor(
 		udpTimeout:             options.UDPTimeout,
 		handler:                options.Handler,
 		logger:                 options.Logger,
+		pauseManager:           pause.ManagerFromContext(options.Context),
 	}
 	return gStack, nil
 }
@@ -76,6 +79,10 @@ func (t *GVisor) Start() error {
 		return err
 	}
 	tcpForwarder := tcp.NewForwarder(ipStack, 0, 1024, func(r *tcp.ForwarderRequest) {
+		if t.pauseManager != nil && pause.IsPaused(t.pauseManager.NetworkPauseChan()) {
+			r.Complete(true)
+			return
+		}
 		var wq waiter.Queue
 		handshakeCtx, cancel := context.WithCancel(context.Background())
 		go func() {
@@ -117,6 +124,10 @@ func (t *GVisor) Start() error {
 	ipStack.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
 	if !t.endpointIndependentNat {
 		udpForwarder := udp.NewForwarder(ipStack, func(request *udp.ForwarderRequest) {
+			if t.pauseManager != nil && pause.IsPaused(t.pauseManager.NetworkPauseChan()) {
+				(*gRequest)(unsafe.Pointer(request)).pkt.DecRef()
+				return
+			}
 			var wq waiter.Queue
 			endpoint, err := request.CreateEndpoint(&wq)
 			if err != nil {
@@ -143,7 +154,7 @@ func (t *GVisor) Start() error {
 		})
 		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, udpForwarder.HandlePacket)
 	} else {
-		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, NewUDPForwarder(t.ctx, ipStack, t.handler, t.udpTimeout).HandlePacket)
+		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, NewUDPForwarder(t.ctx, ipStack, t.handler, t.udpTimeout, t.pauseManager).HandlePacket)
 	}
 
 	t.stack = ipStack

stack_gvisor_udp.go

@@ -22,24 +22,27 @@ import (
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
 	"github.com/sagernet/sing/common/udpnat"
+	"github.com/sagernet/sing/service/pause"
 )
 
 type UDPForwarder struct {
-	ctx    context.Context
-	stack  *stack.Stack
-	udpNat *udpnat.Service[netip.AddrPort]
+	ctx          context.Context
+	stack        *stack.Stack
+	udpNat       *udpnat.Service[netip.AddrPort]
+	pauseManager pause.Manager
 
 	// cache
 	cacheProto  tcpip.NetworkProtocolNumber
 	cacheID     stack.TransportEndpointID
 	cachePacket stack.PacketBufferPtr
 }
 
-func NewUDPForwarder(ctx context.Context, stack *stack.Stack, handler Handler, udpTimeout int64) *UDPForwarder {
+func NewUDPForwarder(ctx context.Context, stack *stack.Stack, handler Handler, udpTimeout int64, pauseManager pause.Manager) *UDPForwarder {
 	return &UDPForwarder{
-		ctx:    ctx,
-		stack:  stack,
-		udpNat: udpnat.New[netip.AddrPort](udpTimeout, handler),
+		ctx:          ctx,
+		stack:        stack,
+		udpNat:       udpnat.New[netip.AddrPort](udpTimeout, handler),
+		pauseManager: pauseManager,
 	}
 }
 
@@ -59,7 +62,7 @@ func (f *UDPForwarder) HandlePacket(id stack.TransportEndpointID, pkt stack.Pack
 	})
 	f.cacheID = id
 	f.cachePacket = pkt
-	f.udpNat.NewPacket(
+	f.udpNat.NewContextPacket1(
 		f.ctx,
 		upstreamMetadata.Source.AddrPort(),
 		sBuffer,
@@ -69,14 +72,17 @@ func (f *UDPForwarder) HandlePacket(id stack.TransportEndpointID, pkt stack.Pack
 	return true
 }
 
-func (f *UDPForwarder) newUDPConn(natConn N.PacketConn) N.PacketWriter {
-	return &UDPBackWriter{
+func (f *UDPForwarder) newUDPConn(natConn N.PacketConn) (context.Context, N.PacketWriter, error) {
+	if f.pauseManager != nil && pause.IsPaused(f.pauseManager.NetworkPauseChan()) {
+		return nil, nil, pause.ErrPaused
+	}
+	return f.ctx, &UDPBackWriter{
 		stack:         f.stack,
 		source:        f.cacheID.RemoteAddress,
 		sourcePort:    f.cacheID.RemotePort,
 		sourceNetwork: f.cacheProto,
 		packet:        f.cachePacket.IncRef(),
-	}
+	}, nil
 }
 
 type UDPBackWriter struct {

stack_mixed.go

@@ -82,7 +82,7 @@ func (m *Mixed) Start() error {
 		})
 		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, udpForwarder.HandlePacket)
 	} else {
-		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, NewUDPForwarder(m.ctx, ipStack, m.handler, m.udpTimeout).HandlePacket)
+		ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, NewUDPForwarder(m.ctx, ipStack, m.handler, m.udpTimeout, m.pauseManager).HandlePacket)
 	}
 	m.stack = ipStack
 	m.endpoint = endpoint

stack_system.go

@@ -16,6 +16,7 @@ import (
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
 	"github.com/sagernet/sing/common/udpnat"
+	"github.com/sagernet/sing/service/pause"
 )
 
 type System struct {
@@ -40,6 +41,7 @@ type System struct {
 	udpNat             *udpnat.Service[netip.AddrPort]
 	bindInterface      bool
 	interfaceFinder    control.InterfaceFinder
+	pauseManager       pause.Manager
 }
 
 type Session struct {
@@ -62,6 +64,7 @@ func NewSystem(options StackOptions) (Stack, error) {
 		inet6Prefixes:   options.Inet6Address,
 		bindInterface:   options.ForwarderBindInterface,
 		interfaceFinder: options.InterfaceFinder,
+		pauseManager:    pause.ManagerFromContext(options.Context),
 	}
 	if len(options.Inet4Address) > 0 {
 		if options.Inet4Address[0].Bits() == 32 {
@@ -197,6 +200,11 @@ func (s *System) acceptLoop(listener net.Listener) {
 		if err != nil {
 			return
 		}
+		if s.pauseManager != nil && pause.IsPaused(s.pauseManager.NetworkPauseChan()) {
+			conn.(*net.TCPConn).SetLinger(0)
+			conn.Close()
+			continue
+		}
 		connPort := M.SocksaddrFromNet(conn.RemoteAddr()).Port
 		session := s.tcpNat.LookupBack(connPort)
 		if session == nil {
@@ -330,11 +338,14 @@ func (s *System) processIPv4UDP(packet clashtcpip.IPv4Packet, header clashtcpip.
 		Source:      M.SocksaddrFromNetIP(source),
 		Destination: M.SocksaddrFromNetIP(destination),
 	}
-	s.udpNat.NewPacket(s.ctx, source, data.ToOwned(), metadata, func(natConn N.PacketConn) N.PacketWriter {
+	s.udpNat.NewContextPacket1(s.ctx, source, data.ToOwned(), metadata, func(natConn N.PacketConn) (context.Context, N.PacketWriter, error) {
+		if s.pauseManager != nil && pause.IsPaused(s.pauseManager.NetworkPauseChan()) {
+			return nil, nil, pause.ErrPaused
+		}
 		headerLen := packet.HeaderLen() + clashtcpip.UDPHeaderSize
 		headerCopy := make([]byte, headerLen)
 		copy(headerCopy, packet[:headerLen])
-		return &systemUDPPacketWriter4{s.tun, headerCopy, source}
+		return s.ctx, &systemUDPPacketWriter4{s.tun, headerCopy, source}, nil
 	})
 	return nil
 }
@@ -353,11 +364,14 @@ func (s *System) processIPv6UDP(packet clashtcpip.IPv6Packet, header clashtcpip.
 		Source:      M.SocksaddrFromNetIP(source),
 		Destination: M.SocksaddrFromNetIP(destination),
 	}
-	s.udpNat.NewPacket(s.ctx, source, data.ToOwned(), metadata, func(natConn N.PacketConn) N.PacketWriter {
+	s.udpNat.NewContextPacket1(s.ctx, source, data.ToOwned(), metadata, func(natConn N.PacketConn) (context.Context, N.PacketWriter, error) {
+		if s.pauseManager != nil && pause.IsPaused(s.pauseManager.NetworkPauseChan()) {
+			return nil, nil, pause.ErrPaused
+		}
 		headerLen := len(packet) - int(header.Length()) + clashtcpip.UDPHeaderSize
 		headerCopy := make([]byte, headerLen)
 		copy(headerCopy, packet[:headerLen])
-		return &systemUDPPacketWriter6{s.tun, headerCopy, source}
+		return s.ctx, &systemUDPPacketWriter6{s.tun, headerCopy, source}, nil
 	})
 	return nil
 }