diff --git a/common/udpnat/service.go b/common/udpnat/service.go index d6cf97b..91cabf0 100644 --- a/common/udpnat/service.go +++ b/common/udpnat/service.go @@ -5,6 +5,7 @@ import ( "io" "net" "os" + "sync/atomic" "time" "github.com/sagernet/sing/common" @@ -113,25 +114,40 @@ type packet struct { } type conn struct { - ctx context.Context - cancel context.CancelFunc - data chan packet - localAddr M.Socksaddr - remoteAddr M.Socksaddr - source N.PacketWriter - fastClose bool + ctx context.Context + cancel context.CancelFunc + data chan packet + localAddr M.Socksaddr + remoteAddr M.Socksaddr + source N.PacketWriter + fastClose bool + readDeadline atomic.Value } func (c *conn) ReadPacketThreadSafe() (buffer *buf.Buffer, addr M.Socksaddr, err error) { + var deadline <-chan time.Time + if d, ok := c.readDeadline.Load().(time.Time); ok && !d.IsZero() { + timer := time.NewTimer(time.Until(d)) + defer timer.Stop() + deadline = timer.C + } select { case p := <-c.data: return p.data, p.destination, nil case <-c.ctx.Done(): - return nil, M.Socksaddr{}, c.ctx.Err() + return nil, M.Socksaddr{}, io.ErrClosedPipe + case <-deadline: + return nil, M.Socksaddr{}, os.ErrDeadlineExceeded } } func (c *conn) ReadPacket(buffer *buf.Buffer) (addr M.Socksaddr, err error) { + var deadline <-chan time.Time + if d, ok := c.readDeadline.Load().(time.Time); ok && !d.IsZero() { + timer := time.NewTimer(time.Until(d)) + defer timer.Stop() + deadline = timer.C + } select { case p := <-c.data: _, err = buffer.ReadFrom(p.data) @@ -139,6 +155,8 @@ func (c *conn) ReadPacket(buffer *buf.Buffer) (addr M.Socksaddr, err error) { return p.destination, err case <-c.ctx.Done(): return M.Socksaddr{}, io.ErrClosedPipe + case <-deadline: + return M.Socksaddr{}, os.ErrDeadlineExceeded } } @@ -191,7 +209,8 @@ func (c *conn) SetDeadline(t time.Time) error { } func (c *conn) SetReadDeadline(t time.Time) error { - return os.ErrInvalid + c.readDeadline.Store(t) + return nil } func (c *conn) SetWriteDeadline(t time.Time) error { diff --git a/go.mod b/go.mod index fd587ee..d6bd075 100644 --- a/go.mod +++ b/go.mod @@ -2,4 +2,4 @@ module github.com/sagernet/sing go 1.18 -require golang.org/x/sys v0.0.0-20220615213510-4f61da869c0c +require golang.org/x/sys v0.0.0-20220627191245-f75cf1eec38b diff --git a/go.sum b/go.sum index f140671..a2987a2 100644 --- a/go.sum +++ b/go.sum @@ -1,2 +1,2 @@ -golang.org/x/sys v0.0.0-20220615213510-4f61da869c0c h1:aFV+BgZ4svzjfabn8ERpuB4JI4N6/rdy1iusx77G3oU= -golang.org/x/sys v0.0.0-20220615213510-4f61da869c0c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220627191245-f75cf1eec38b h1:2n253B2r0pYSmEV+UNCQoPfU/FiaizQEK5Gu4Bq4JE8= +golang.org/x/sys v0.0.0-20220627191245-f75cf1eec38b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=