From ccfe5c0f0f2bfe9c05e13d0c28ca47a2bcd5ca8e Mon Sep 17 00:00:00 2001 From: wwqgtxx Date: Sun, 24 Aug 2025 18:29:39 +0800 Subject: [PATCH] ping: Rewrite UnprivilegedConn --- ping/socket_linux_unprivileged.go | 141 ++++++++++++++++++------------ 1 file changed, 85 insertions(+), 56 deletions(-) diff --git a/ping/socket_linux_unprivileged.go b/ping/socket_linux_unprivileged.go index 1da0c83..78d1026 100644 --- a/ping/socket_linux_unprivileged.go +++ b/ping/socket_linux_unprivileged.go @@ -5,14 +5,16 @@ import ( "net" "net/netip" "os" + "sync" "time" "github.com/sagernet/sing-tun/internal/gtcpip/checksum" "github.com/sagernet/sing-tun/internal/gtcpip/header" - "github.com/sagernet/sing/common/atomic" + "github.com/sagernet/sing/common" "github.com/sagernet/sing/common/buf" "github.com/sagernet/sing/common/control" M "github.com/sagernet/sing/common/metadata" + "github.com/sagernet/sing/common/pipe" ) type UnprivilegedConn struct { @@ -21,7 +23,9 @@ type UnprivilegedConn struct { controlFunc control.Func destination netip.Addr receiveChan chan *unprivilegedResponse - readDeadline atomic.TypedValue[time.Time] + readDeadline pipe.Deadline + natMap map[uint16]net.Conn + natMapMutex sync.Mutex } type unprivilegedResponse struct { @@ -38,11 +42,13 @@ func newUnprivilegedConn(ctx context.Context, controlFunc control.Func, destinat conn.Close() ctx, cancel := context.WithCancel(ctx) return &UnprivilegedConn{ - ctx: ctx, - cancel: cancel, - controlFunc: controlFunc, - destination: destination, - receiveChan: make(chan *unprivilegedResponse), + ctx: ctx, + cancel: cancel, + controlFunc: controlFunc, + destination: destination, + receiveChan: make(chan *unprivilegedResponse), + readDeadline: pipe.MakeDeadline(), + natMap: make(map[uint16]net.Conn), }, nil } @@ -55,6 +61,8 @@ func (c *UnprivilegedConn) Read(b []byte) (n int, err error) { return case <-c.ctx.Done(): return 0, os.ErrClosed + case <-c.readDeadline.Wait(): + return 0, os.ErrDeadlineExceeded } } @@ -69,14 +77,12 @@ func (c *UnprivilegedConn) ReadMsg(b []byte, oob []byte) (n, oobn int, addr neti return case <-c.ctx.Done(): return 0, 0, netip.Addr{}, os.ErrClosed + case <-c.readDeadline.Wait(): + return 0, 0, netip.Addr{}, os.ErrDeadlineExceeded } } func (c *UnprivilegedConn) Write(b []byte) (n int, err error) { - conn, err := connect(false, c.controlFunc, c.destination) - if err != nil { - return - } var identifier uint16 if !c.destination.Is6() { icmpHdr := header.ICMPv4(b) @@ -85,62 +91,85 @@ func (c *UnprivilegedConn) Write(b []byte) (n int, err error) { icmpHdr := header.ICMPv6(b) identifier = icmpHdr.Ident() } - if readDeadline := c.readDeadline.Load(); !readDeadline.IsZero() { - conn.SetReadDeadline(readDeadline) + + c.natMapMutex.Lock() + if err = c.ctx.Err(); err != nil { + c.natMapMutex.Unlock() + return 0, err } + conn, ok := c.natMap[identifier] + if !ok { + conn, err = connect(false, c.controlFunc, c.destination) + if err != nil { + c.natMapMutex.Unlock() + return 0, err + } + go c.fetchResponse(conn.(*net.UDPConn), identifier) + } + c.natMapMutex.Unlock() + n, err = conn.Write(b) if err != nil { - conn.Close() + c.removeConn(conn.(*net.UDPConn), identifier) return } - go c.fetchResponse(conn, identifier) return } -func (c *UnprivilegedConn) fetchResponse(conn net.Conn, identifier uint16) { - done := make(chan struct{}) - defer close(done) - go func() { - select { - case <-c.ctx.Done(): - case <-done: +func (c *UnprivilegedConn) fetchResponse(conn *net.UDPConn, identifier uint16) { + defer c.removeConn(conn, identifier) + for { + buffer := buf.NewPacket() + cmsgBuffer := buf.NewSize(1024) + n, oobN, _, addr, err := conn.ReadMsgUDPAddrPort(buffer.FreeBytes(), cmsgBuffer.FreeBytes()) + if err != nil { + buffer.Release() + cmsgBuffer.Release() + return + } + buffer.Truncate(n) + cmsgBuffer.Truncate(oobN) + if !c.destination.Is6() { + icmpHdr := header.ICMPv4(buffer.Bytes()) + icmpHdr.SetIdent(identifier) + icmpHdr.SetChecksum(0) + icmpHdr.SetChecksum(header.ICMPv4Checksum(icmpHdr[:header.ICMPv4MinimumSize], checksum.Checksum(icmpHdr.Payload(), 0))) + } else { + icmpHdr := header.ICMPv6(buffer.Bytes()) + icmpHdr.SetIdent(identifier) + // offload checksum here since we don't have source address here + } + select { + case c.receiveChan <- &unprivilegedResponse{ + Buffer: buffer, + Cmsg: cmsgBuffer, + Addr: addr.Addr(), + }: + case <-c.ctx.Done(): + buffer.Release() + cmsgBuffer.Release() + return } - conn.Close() - }() - buffer := buf.NewPacket() - cmsgBuffer := buf.NewSize(1024) - n, oobN, _, addr, err := conn.(*net.UDPConn).ReadMsgUDPAddrPort(buffer.FreeBytes(), cmsgBuffer.FreeBytes()) - if err != nil { - buffer.Release() - cmsgBuffer.Release() - return - } - buffer.Truncate(n) - cmsgBuffer.Truncate(oobN) - if !c.destination.Is6() { - icmpHdr := header.ICMPv4(buffer.Bytes()) - icmpHdr.SetIdent(identifier) - icmpHdr.SetChecksum(0) - icmpHdr.SetChecksum(header.ICMPv4Checksum(icmpHdr[:header.ICMPv4MinimumSize], checksum.Checksum(icmpHdr.Payload(), 0))) - } else { - icmpHdr := header.ICMPv6(buffer.Bytes()) - icmpHdr.SetIdent(identifier) - // offload checksum here since we don't have source address here - } - select { - case c.receiveChan <- &unprivilegedResponse{ - Buffer: buffer, - Cmsg: cmsgBuffer, - Addr: addr.Addr(), - }: - case <-c.ctx.Done(): - buffer.Release() - cmsgBuffer.Release() } } +func (c *UnprivilegedConn) removeConn(conn *net.UDPConn, identifier uint16) { + c.natMapMutex.Lock() + _ = conn.Close() + if c.natMap[identifier] == conn { + delete(c.natMap, identifier) + } + c.natMapMutex.Unlock() +} + func (c *UnprivilegedConn) Close() error { + c.natMapMutex.Lock() c.cancel() + for _, conn := range c.natMap { + _ = conn.Close() + } + common.ClearMap(c.natMap) + c.natMapMutex.Unlock() return nil } @@ -153,14 +182,14 @@ func (c *UnprivilegedConn) RemoteAddr() net.Addr { } func (c *UnprivilegedConn) SetDeadline(t time.Time) error { - return os.ErrInvalid + return c.SetReadDeadline(t) } func (c *UnprivilegedConn) SetReadDeadline(t time.Time) error { - c.readDeadline.Store(t) + c.readDeadline.Set(t) return nil } func (c *UnprivilegedConn) SetWriteDeadline(t time.Time) error { - return os.ErrInvalid + return nil }