Add flow dispatcher

ping/destination.go

@@ -9,7 +9,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/sagernet/sing-tun"
 	"github.com/sagernet/sing-tun/gtcpip/header"
 	"github.com/sagernet/sing/common/buf"
 	"github.com/sagernet/sing/common/control"
@@ -20,14 +19,16 @@ import (
 // Although its theoretical maximum may be 64k, I don’t yet know of any practical use case for that. For memory-usage reasons, I’m just using a 2k buffer.
 const maxICMPPacketSize = 2048
 
-var _ tun.DirectRouteDestination = (*Destination)(nil)
+type PacketWriter interface {
+	WritePacket(packet []byte) error
+}
 
 type Destination struct {
 	conn          *Conn
 	ctx           context.Context
 	logger        logger.ContextLogger
 	destination   netip.Addr
-	routeContext  tun.DirectRouteContext
+	writer        PacketWriter
 	timeout       time.Duration
 	requestAccess sync.Mutex
 	requests      map[pingRequest]time.Time
@@ -45,9 +46,9 @@ func ConnectDestination(
 	logger logger.ContextLogger,
 	controlFunc control.Func,
 	destination netip.Addr,
-	routeContext tun.DirectRouteContext,
+	writer PacketWriter,
 	timeout time.Duration,
-) (tun.DirectRouteDestination, error) {
+) (*Destination, error) {
 	var (
 		conn *Conn
 		err  error
@@ -65,13 +66,13 @@ func ConnectDestination(
 		return nil, err
 	}
 	d := &Destination{
-		conn:         conn,
-		ctx:          ctx,
-		logger:       logger,
-		destination:  destination,
-		routeContext: routeContext,
-		timeout:      timeout,
-		requests:     make(map[pingRequest]time.Time),
+		conn:        conn,
+		ctx:         ctx,
+		logger:      logger,
+		destination: destination,
+		writer:      writer,
+		timeout:     timeout,
+		requests:    make(map[pingRequest]time.Time),
 	}
 	go d.loopRead()
 	return d, nil
@@ -158,7 +159,7 @@ func (d *Destination) loopRead() {
 			}
 			d.logger.TraceContext(d.ctx, "read ICMPv6 echo reply from ", ipHdr.SourceAddr(), " to ", ipHdr.DestinationAddr(), " id ", icmpHdr.Ident(), " seq ", icmpHdr.Sequence())
 		}
-		err = d.routeContext.WritePacket(buffer.Bytes())
+		err = d.writer.WritePacket(buffer.Bytes())
 		if err != nil {
 			d.logger.ErrorContext(d.ctx, E.Cause(err, "write ICMP echo reply"))
 		}

ping/destination_gvisor.go

@@ -1,143 +0,0 @@
-//go:build with_gvisor
-
-package ping
-
-import (
-	"context"
-	"net/netip"
-	"time"
-
-	"github.com/sagernet/gvisor/pkg/tcpip"
-	"github.com/sagernet/gvisor/pkg/tcpip/adapters/gonet"
-	"github.com/sagernet/gvisor/pkg/tcpip/header"
-	"github.com/sagernet/gvisor/pkg/tcpip/stack"
-	"github.com/sagernet/gvisor/pkg/tcpip/transport"
-	"github.com/sagernet/gvisor/pkg/waiter"
-	"github.com/sagernet/sing-tun"
-	"github.com/sagernet/sing/common"
-	"github.com/sagernet/sing/common/buf"
-	E "github.com/sagernet/sing/common/exceptions"
-	"github.com/sagernet/sing/common/logger"
-)
-
-var _ tun.DirectRouteDestination = (*GVisorDestination)(nil)
-
-type GVisorDestination struct {
-	ctx        context.Context
-	logger     logger.ContextLogger
-	endpoint   tcpip.Endpoint
-	conn       *gonet.TCPConn
-	rewriter   *SourceRewriter
-	timeout    time.Duration
-	lastActive common.TypedValue[time.Time]
-}
-
-func ConnectGVisor(
-	ctx context.Context, logger logger.ContextLogger,
-	sourceAddress, destinationAddress netip.Addr,
-	routeContext tun.DirectRouteContext,
-	stack *stack.Stack,
-	bindAddress4, bindAddress6 netip.Addr,
-	timeout time.Duration,
-) (*GVisorDestination, error) {
-	var (
-		bindAddress tcpip.Address
-		wq          waiter.Queue
-		endpoint    tcpip.Endpoint
-		gErr        tcpip.Error
-	)
-	if !destinationAddress.Is6() {
-		if !bindAddress4.IsValid() {
-			return nil, E.New("missing IPv4 interface address")
-		}
-		bindAddress = tun.AddressFromAddr(bindAddress4)
-		endpoint, gErr = stack.NewRawEndpoint(header.ICMPv4ProtocolNumber, header.IPv4ProtocolNumber, &wq, true)
-	} else {
-		if !bindAddress6.IsValid() {
-			return nil, E.New("missing IPv6 interface address")
-		}
-		bindAddress = tun.AddressFromAddr(bindAddress6)
-		endpoint, gErr = stack.NewRawEndpoint(header.ICMPv6ProtocolNumber, header.IPv6ProtocolNumber, &wq, true)
-	}
-	if gErr != nil {
-		return nil, gonet.TranslateNetstackError(gErr)
-	}
-	gErr = endpoint.Bind(tcpip.FullAddress{
-		NIC:  1,
-		Addr: bindAddress,
-	})
-	if gErr != nil {
-		return nil, gonet.TranslateNetstackError(gErr)
-	}
-	gErr = endpoint.Connect(tcpip.FullAddress{
-		NIC:  1,
-		Addr: tun.AddressFromAddr(destinationAddress),
-	})
-	if gErr != nil {
-		return nil, gonet.TranslateNetstackError(gErr)
-	}
-	endpoint.SocketOptions().SetHeaderIncluded(true)
-	rewriter := NewSourceRewriter(ctx, logger, bindAddress4, bindAddress6)
-	rewriter.CreateSession(tun.DirectRouteSession{Source: sourceAddress, Destination: destinationAddress}, routeContext)
-	destination := &GVisorDestination{
-		ctx:      ctx,
-		logger:   logger,
-		endpoint: endpoint,
-		conn:     gonet.NewTCPConn(&wq, endpoint),
-		rewriter: rewriter,
-		timeout:  timeout,
-	}
-	destination.lastActive.Store(time.Now())
-	go destination.loopRead()
-	return destination, nil
-}
-
-func (d *GVisorDestination) loopRead() {
-	defer d.endpoint.Close()
-	for {
-		deadline := d.lastActive.Load().Add(d.timeout)
-		if !time.Now().Before(deadline) {
-			return
-		}
-		err := d.conn.SetReadDeadline(deadline)
-		if err != nil {
-			d.logger.ErrorContext(d.ctx, E.Cause(err, "set read deadline for ICMP conn"))
-		}
-		buffer := buf.NewSize(maxICMPPacketSize)
-		n, err := d.conn.Read(buffer.FreeBytes())
-		if err != nil {
-			buffer.Release()
-			if E.IsTimeout(err) {
-				continue
-			}
-			if !E.IsClosed(err) {
-				d.logger.ErrorContext(d.ctx, E.Cause(err, "receive ICMP echo reply"))
-			}
-			return
-		}
-		buffer.Truncate(n)
-		var matched bool
-		matched, err = d.rewriter.WriteBack(buffer.Bytes())
-		if err != nil {
-			d.logger.ErrorContext(d.ctx, E.Cause(err, "write ICMP echo reply"))
-		}
-		if matched {
-			d.lastActive.Store(time.Now())
-		}
-		buffer.Release()
-	}
-}
-
-func (d *GVisorDestination) WritePacket(packet *buf.Buffer) error {
-	d.lastActive.Store(time.Now())
-	d.rewriter.RewritePacket(packet.Bytes())
-	return common.Error(d.conn.Write(packet.Bytes()))
-}
-
-func (d *GVisorDestination) Close() error {
-	return d.conn.Close()
-}
-
-func (d *GVisorDestination) IsClosed() bool {
-	return transport.DatagramEndpointState(d.endpoint.State()) == transport.DatagramEndpointStateClosed
-}

ping/destination_rewriter.go

@@ -1,79 +0,0 @@
-package ping
-
-import (
-	"net/netip"
-
-	"github.com/sagernet/sing-tun"
-	"github.com/sagernet/sing-tun/gtcpip/header"
-	"github.com/sagernet/sing/common/buf"
-)
-
-type DestinationWriter struct {
-	tun.DirectRouteDestination
-	destination netip.Addr
-}
-
-func NewDestinationWriter(routeDestination tun.DirectRouteDestination, destination netip.Addr) *DestinationWriter {
-	return &DestinationWriter{routeDestination, destination}
-}
-
-func (w *DestinationWriter) WritePacket(packet *buf.Buffer) error {
-	var ipHdr header.Network
-	switch header.IPVersion(packet.Bytes()) {
-	case header.IPv4Version:
-		ipHdr = header.IPv4(packet.Bytes())
-	case header.IPv6Version:
-		ipHdr = header.IPv6(packet.Bytes())
-	default:
-		return w.DirectRouteDestination.WritePacket(packet)
-	}
-	ipHdr.SetDestinationAddr(w.destination)
-	if ipHdr4, isIPv4 := ipHdr.(header.IPv4); isIPv4 {
-		ipHdr4.SetChecksum(^ipHdr4.CalculateChecksum())
-	}
-	if ipHdr.TransportProtocol() == header.ICMPv6ProtocolNumber {
-		icmpHdr := header.ICMPv6(ipHdr.Payload())
-		icmpHdr.SetChecksum(header.ICMPv6Checksum(header.ICMPv6ChecksumParams{
-			Header: icmpHdr,
-			Src:    ipHdr.SourceAddressSlice(),
-			Dst:    ipHdr.DestinationAddressSlice(),
-		}))
-	}
-	return w.DirectRouteDestination.WritePacket(packet)
-}
-
-type ContextDestinationWriter struct {
-	tun.DirectRouteContext
-	destination netip.Addr
-}
-
-func NewContextDestinationWriter(context tun.DirectRouteContext, destination netip.Addr) *ContextDestinationWriter {
-	return &ContextDestinationWriter{
-		context, destination,
-	}
-}
-
-func (w *ContextDestinationWriter) WritePacket(packet []byte) error {
-	var ipHdr header.Network
-	switch header.IPVersion(packet) {
-	case header.IPv4Version:
-		ipHdr = header.IPv4(packet)
-	case header.IPv6Version:
-		ipHdr = header.IPv6(packet)
-	default:
-		return w.DirectRouteContext.WritePacket(packet)
-	}
-	ipHdr.SetSourceAddr(w.destination)
-	if ipHdr4, isIPv4 := ipHdr.(header.IPv4); isIPv4 {
-		ipHdr4.SetChecksum(^ipHdr4.CalculateChecksum())
-	}
-	if ipHdr.TransportProtocol() == header.ICMPv6ProtocolNumber {
-		icmpHdr := header.ICMPv6(ipHdr.Payload())
-		icmpHdr.SetChecksum(header.ICMPv6Checksum(header.ICMPv6ChecksumParams{
-			Header: icmpHdr,
-			Src:    ipHdr.SourceAddressSlice(),
-			Dst:    ipHdr.DestinationAddressSlice(),
-		}))
-	}
-	return w.DirectRouteContext.WritePacket(packet)
-}

ping/port.go

@@ -0,0 +1,194 @@
+package ping
+
+import (
+	"context"
+	"net/netip"
+	"slices"
+	"sync"
+	"time"
+
+	"github.com/sagernet/sing-tun"
+	"github.com/sagernet/sing-tun/gtcpip/header"
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/control"
+	E "github.com/sagernet/sing/common/exceptions"
+	"github.com/sagernet/sing/common/logger"
+)
+
+const defaultFlowTimeout = time.Minute
+
+type Port struct {
+	ctx         context.Context
+	logger      logger.ContextLogger
+	controlFunc func(destination netip.Addr) control.Func
+	timeout     time.Duration
+
+	returnAccess sync.Mutex
+	returnPaths  []tun.Return
+
+	flowAccess sync.Mutex
+	flows      map[flowKey]*Destination
+	lastSweep  time.Time
+}
+
+type flowKey struct {
+	source      netip.Addr
+	destination netip.Addr
+	identifier  uint16
+}
+
+func NewPort(ctx context.Context, logger logger.ContextLogger, controlFunc func(destination netip.Addr) control.Func, timeout time.Duration) *Port {
+	if timeout <= 0 {
+		timeout = defaultFlowTimeout
+	}
+	return &Port{
+		ctx:         ctx,
+		logger:      logger,
+		controlFunc: controlFunc,
+		timeout:     timeout,
+		flows:       make(map[flowKey]*Destination),
+	}
+}
+
+func (p *Port) PortAddresses() (netip.Addr, netip.Addr) {
+	return netip.IPv4Unspecified(), netip.IPv6Unspecified()
+}
+
+func (p *Port) PortMTU() uint32 {
+	return 0
+}
+
+func (p *Port) AttachReturn(returnPath tun.Return) error {
+	p.returnAccess.Lock()
+	defer p.returnAccess.Unlock()
+	if slices.Contains(p.returnPaths, returnPath) {
+		return nil
+	}
+	p.returnPaths = append(p.returnPaths[:len(p.returnPaths):len(p.returnPaths)], returnPath)
+	return nil
+}
+
+func (p *Port) DetachReturn(returnPath tun.Return) error {
+	p.returnAccess.Lock()
+	defer p.returnAccess.Unlock()
+	returnPaths := make([]tun.Return, 0, len(p.returnPaths))
+	for _, existing := range p.returnPaths {
+		if existing != returnPath {
+			returnPaths = append(returnPaths, existing)
+		}
+	}
+	p.returnPaths = returnPaths
+	return nil
+}
+
+func (p *Port) WritePackets(packets [][]byte) error {
+	var errs []error
+	for _, packet := range packets {
+		err := p.writePacket(packet)
+		if err != nil {
+			errs = append(errs, err)
+		}
+	}
+	return E.Errors(errs...)
+}
+
+func (p *Port) writePacket(packet []byte) error {
+	var (
+		source      netip.Addr
+		destination netip.Addr
+		identifier  uint16
+	)
+	switch header.IPVersion(packet) {
+	case header.IPv4Version:
+		ipHdr := header.IPv4(packet)
+		if !ipHdr.IsValid(len(packet)) || ipHdr.TransportProtocol() != header.ICMPv4ProtocolNumber || ipHdr.PayloadLength() < header.ICMPv4MinimumSize {
+			return nil
+		}
+		icmpHdr := header.ICMPv4(ipHdr.Payload())
+		if icmpHdr.Type() != header.ICMPv4Echo || icmpHdr.Code() != 0 {
+			return nil
+		}
+		source = ipHdr.SourceAddr()
+		destination = ipHdr.DestinationAddr()
+		identifier = icmpHdr.Ident()
+	case header.IPv6Version:
+		ipHdr := header.IPv6(packet)
+		if !ipHdr.IsValid(len(packet)) || ipHdr.TransportProtocol() != header.ICMPv6ProtocolNumber || ipHdr.PayloadLength() < header.ICMPv6MinimumSize {
+			return nil
+		}
+		icmpHdr := header.ICMPv6(ipHdr.Payload())
+		if icmpHdr.Type() != header.ICMPv6EchoRequest || icmpHdr.Code() != 0 {
+			return nil
+		}
+		source = ipHdr.SourceAddr()
+		destination = ipHdr.DestinationAddr()
+		identifier = icmpHdr.Ident()
+	default:
+		return nil
+	}
+	flow, err := p.flowFor(source, destination, identifier)
+	if err != nil {
+		return E.Cause(err, "connect ICMP flow to ", destination)
+	}
+	return flow.WritePacket(buf.As(packet))
+}
+
+func (p *Port) flowFor(source netip.Addr, destination netip.Addr, identifier uint16) (*Destination, error) {
+	key := flowKey{source: source, destination: destination, identifier: identifier}
+	p.flowAccess.Lock()
+	defer p.flowAccess.Unlock()
+	now := time.Now()
+	if now.Sub(p.lastSweep) >= p.timeout {
+		p.lastSweep = now
+		for oldKey, oldFlow := range p.flows {
+			if oldFlow.IsClosed() {
+				delete(p.flows, oldKey)
+			}
+		}
+	}
+	flow, loaded := p.flows[key]
+	if loaded && !flow.IsClosed() {
+		return flow, nil
+	}
+	var controlFunc control.Func
+	if p.controlFunc != nil {
+		controlFunc = p.controlFunc(destination)
+	}
+	flow, err := ConnectDestination(p.ctx, p.logger, controlFunc, destination, portWriter{p}, p.timeout)
+	if err != nil {
+		return nil, err
+	}
+	p.flows[key] = flow
+	return flow, nil
+}
+
+type portWriter struct {
+	port *Port
+}
+
+func (w portWriter) WritePacket(packet []byte) error {
+	w.port.returnAccess.Lock()
+	returnPaths := w.port.returnPaths
+	w.port.returnAccess.Unlock()
+	for _, returnPath := range returnPaths {
+		headroom := returnPath.ReturnHeadroom()
+		buffer := make([]byte, headroom+len(packet))
+		copy(buffer[headroom:], packet)
+		unconsumed := returnPath.ReturnPackets([][]byte{buffer})
+		if len(unconsumed) == 0 {
+			return nil
+		}
+	}
+	return nil
+}
+
+func (p *Port) Close() error {
+	p.flowAccess.Lock()
+	defer p.flowAccess.Unlock()
+	var errs []error
+	for key, flow := range p.flows {
+		errs = append(errs, flow.Close())
+		delete(p.flows, key)
+	}
+	return E.Errors(errs...)
+}

ping/source_rewriter.go

@@ -1,150 +0,0 @@
-package ping
-
-import (
-	"context"
-	"net/netip"
-	"sync"
-
-	"github.com/sagernet/sing-tun"
-	"github.com/sagernet/sing-tun/gtcpip/header"
-	"github.com/sagernet/sing/common/logger"
-)
-
-type SourceRewriter struct {
-	ctx           context.Context
-	logger        logger.ContextLogger
-	access        sync.RWMutex
-	sessions      map[tun.DirectRouteSession]tun.DirectRouteContext
-	sourceAddress map[uint16]netip.Addr
-	inet4Address  netip.Addr
-	inet6Address  netip.Addr
-}
-
-func NewSourceRewriter(ctx context.Context, logger logger.ContextLogger, inet4Address netip.Addr, inet6Address netip.Addr) *SourceRewriter {
-	return &SourceRewriter{
-		ctx:           ctx,
-		logger:        logger,
-		sessions:      make(map[tun.DirectRouteSession]tun.DirectRouteContext),
-		sourceAddress: make(map[uint16]netip.Addr),
-		inet4Address:  inet4Address,
-		inet6Address:  inet6Address,
-	}
-}
-
-func (m *SourceRewriter) CreateSession(session tun.DirectRouteSession, context tun.DirectRouteContext) {
-	m.access.Lock()
-	m.sessions[session] = context
-	m.access.Unlock()
-}
-
-func (m *SourceRewriter) DeleteSession(session tun.DirectRouteSession) {
-	m.access.Lock()
-	delete(m.sessions, session)
-	m.access.Unlock()
-}
-
-func (m *SourceRewriter) RewritePacket(packet []byte) {
-	var ipHdr header.Network
-	var bindAddr netip.Addr
-	switch header.IPVersion(packet) {
-	case header.IPv4Version:
-		ipHdr = header.IPv4(packet)
-		bindAddr = m.inet4Address
-	case header.IPv6Version:
-		ipHdr = header.IPv6(packet)
-		bindAddr = m.inet6Address
-	default:
-		return
-	}
-	sourceAddr := ipHdr.SourceAddr()
-	ipHdr.SetSourceAddr(bindAddr)
-	if ipHdr4, isIPv4 := ipHdr.(header.IPv4); isIPv4 {
-		ipHdr4.SetChecksum(^ipHdr4.CalculateChecksum())
-	}
-	switch ipHdr.TransportProtocol() {
-	case header.ICMPv4ProtocolNumber:
-		icmpHdr := header.ICMPv4(ipHdr.Payload())
-		m.access.Lock()
-		m.sourceAddress[icmpHdr.Ident()] = sourceAddr
-		m.access.Unlock()
-		m.logger.TraceContext(m.ctx, "write ICMPv4 echo request from ", ipHdr.SourceAddr(), " to ", ipHdr.DestinationAddr(), " id ", icmpHdr.Ident(), " seq ", icmpHdr.Sequence())
-	case header.ICMPv6ProtocolNumber:
-		icmpHdr := header.ICMPv6(ipHdr.Payload())
-		icmpHdr.SetChecksum(header.ICMPv6Checksum(header.ICMPv6ChecksumParams{
-			Header: icmpHdr,
-			Src:    ipHdr.SourceAddressSlice(),
-			Dst:    ipHdr.DestinationAddressSlice(),
-		}))
-		m.access.Lock()
-		m.sourceAddress[icmpHdr.Ident()] = sourceAddr
-		m.access.Unlock()
-		m.logger.TraceContext(m.ctx, "write ICMPv6 echo request from ", ipHdr.SourceAddr(), " to ", ipHdr.DestinationAddr(), " id ", icmpHdr.Ident(), " seq ", icmpHdr.Sequence())
-	}
-}
-
-func (m *SourceRewriter) WriteBack(packet []byte) (bool, error) {
-	var ipHdr header.Network
-	var routeSession tun.DirectRouteSession
-	switch header.IPVersion(packet) {
-	case header.IPv4Version:
-		ipHdr = header.IPv4(packet)
-		routeSession.Destination = ipHdr.SourceAddr()
-	case header.IPv6Version:
-		ipHdr = header.IPv6(packet)
-		routeSession.Destination = ipHdr.SourceAddr()
-	default:
-		return false, nil
-	}
-	switch ipHdr.TransportProtocol() {
-	case header.ICMPv4ProtocolNumber:
-		icmpHdr := header.ICMPv4(ipHdr.Payload())
-		m.access.Lock()
-		ident := icmpHdr.Ident()
-		source, loaded := m.sourceAddress[ident]
-		if !loaded {
-			m.access.Unlock()
-			return false, nil
-		}
-		delete(m.sourceAddress, icmpHdr.Ident())
-		m.access.Unlock()
-		routeSession.Source = source
-	case header.ICMPv6ProtocolNumber:
-		icmpHdr := header.ICMPv6(ipHdr.Payload())
-		m.access.Lock()
-		ident := icmpHdr.Ident()
-		source, loaded := m.sourceAddress[ident]
-		if !loaded {
-			m.access.Unlock()
-			return false, nil
-		}
-		delete(m.sourceAddress, icmpHdr.Ident())
-		m.access.Unlock()
-		routeSession.Source = source
-	default:
-		return false, nil
-	}
-	m.access.RLock()
-	context, loaded := m.sessions[routeSession]
-	m.access.RUnlock()
-	if !loaded {
-		return false, nil
-	}
-	ipHdr.SetDestinationAddr(routeSession.Source)
-	if ipHdr4, isIPv4 := ipHdr.(header.IPv4); isIPv4 {
-		ipHdr4.SetChecksum(^ipHdr4.CalculateChecksum())
-	}
-	switch ipHdr.TransportProtocol() {
-	case header.ICMPv4ProtocolNumber:
-		icmpHdr := header.ICMPv4(ipHdr.Payload())
-		m.logger.TraceContext(m.ctx, "read ICMPv4 echo reply from ", ipHdr.SourceAddr(), " to ", ipHdr.DestinationAddr(), " id ", icmpHdr.Ident(), " seq ", icmpHdr.Sequence())
-	case header.ICMPv6ProtocolNumber:
-		icmpHdr := header.ICMPv6(ipHdr.Payload())
-		icmpHdr.SetChecksum(header.ICMPv6Checksum(header.ICMPv6ChecksumParams{
-			Header: icmpHdr,
-			Src:    ipHdr.SourceAddressSlice(),
-			Dst:    ipHdr.DestinationAddressSlice(),
-		}))
-		m.logger.TraceContext(m.ctx, "read ICMPv6 echo reply from ", ipHdr.SourceAddr(), " to ", ipHdr.DestinationAddr(), " id ", icmpHdr.Ident(), " seq ", icmpHdr.Sequence())
-	}
-	return true, context.WritePacket(packet)
-}