Add UDPEgressPool

udp_egress_conn.go

@@ -0,0 +1,124 @@
+package tun
+
+import (
+	"net"
+	"net/netip"
+	"sync"
+	"time"
+
+	"github.com/sagernet/sing/common/buf"
+	E "github.com/sagernet/sing/common/exceptions"
+)
+
+type UDPEgressConn struct {
+	anchor     *net.UDPConn
+	pool       *UDPEgressPool
+	packetChan chan udpEgressConnPacket
+	doneChan   chan struct{}
+	closeOnce  sync.Once
+	readWait   sync.WaitGroup
+}
+
+type udpEgressConnPacket struct {
+	buffer *buf.Buffer
+	source netip.AddrPort
+	err    error
+}
+
+func NewUDPEgressConn(anchor *net.UDPConn, pool *UDPEgressPool) *UDPEgressConn {
+	conn := &UDPEgressConn{
+		anchor:     anchor,
+		pool:       pool,
+		packetChan: make(chan udpEgressConnPacket, 64),
+		doneChan:   make(chan struct{}),
+	}
+	conn.readWait.Add(2)
+	go conn.read(anchor.ReadFromUDPAddrPort)
+	go conn.read(pool.ReceiveEgress)
+	return conn
+}
+
+func (c *UDPEgressConn) read(readPacket func([]byte) (int, netip.AddrPort, error)) {
+	defer c.readWait.Done()
+	for {
+		buffer := buf.NewSize(udpEgressBufferSize)
+		dataLength, source, err := readPacket(buffer.FreeBytes())
+		if err != nil {
+			buffer.Release()
+			if E.IsClosed(err) {
+				return
+			}
+			select {
+			case c.packetChan <- udpEgressConnPacket{err: err}:
+			case <-c.doneChan:
+				return
+			}
+			continue
+		}
+		buffer.Extend(dataLength)
+		select {
+		case c.packetChan <- udpEgressConnPacket{buffer: buffer, source: source}:
+		case <-c.doneChan:
+			buffer.Release()
+			return
+		}
+	}
+}
+
+func (c *UDPEgressConn) ReadFromUDPAddrPort(buffer []byte) (int, netip.AddrPort, error) {
+	select {
+	case packet := <-c.packetChan:
+		if packet.err != nil {
+			return 0, netip.AddrPort{}, packet.err
+		}
+		copied := copy(buffer, packet.buffer.Bytes())
+		packet.buffer.Release()
+		return copied, packet.source, nil
+	case <-c.doneChan:
+		return 0, netip.AddrPort{}, net.ErrClosed
+	}
+}
+
+func (c *UDPEgressConn) WriteToUDPAddrPort(buffer []byte, destination netip.AddrPort) (int, error) {
+	memberConn := c.pool.LookupEgress(destination)
+	if memberConn != nil {
+		return memberConn.WriteToUDPAddrPort(buffer, destination)
+	}
+	return c.anchor.WriteToUDPAddrPort(buffer, destination)
+}
+
+func (c *UDPEgressConn) LocalAddr() net.Addr {
+	return c.anchor.LocalAddr()
+}
+
+func (c *UDPEgressConn) SetDeadline(t time.Time) error {
+	return c.anchor.SetDeadline(t)
+}
+
+func (c *UDPEgressConn) SetReadDeadline(t time.Time) error {
+	return c.anchor.SetReadDeadline(t)
+}
+
+func (c *UDPEgressConn) SetWriteDeadline(t time.Time) error {
+	return c.anchor.SetWriteDeadline(t)
+}
+
+func (c *UDPEgressConn) Close() error {
+	c.closeOnce.Do(func() {
+		close(c.doneChan)
+		c.anchor.Close()
+		c.pool.Close()
+		c.readWait.Wait()
+		for {
+			select {
+			case packet := <-c.packetChan:
+				if packet.buffer != nil {
+					packet.buffer.Release()
+				}
+			default:
+				return
+			}
+		}
+	})
+	return nil
+}