124 lines
2.7 KiB
Go
124 lines
2.7 KiB
Go
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
|
|
}
|