Add flow dispatcher

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...)
+}