diff --git a/internal/netconn/pktpipe.go b/internal/netconn/pktpipe.go index 816809e..5841aa7 100644 --- a/internal/netconn/pktpipe.go +++ b/internal/netconn/pktpipe.go @@ -85,6 +85,11 @@ func (p *packetPipe) ReadFrom(b []byte) (int, net.Addr, error) { select { case <-p.done: return 0, nil, net.ErrClosed + case <-p.peer.done: + if n, ok := p.drain(b); ok { + return n, p.remote, nil + } + return 0, nil, net.ErrClosed case <-changed: continue case <-timeout: @@ -100,6 +105,20 @@ func (p *packetPipe) ReadFrom(b []byte) (int, net.Addr, error) { } } +func (p *packetPipe) drain(b []byte) (int, bool) { + select { + case buf, ok := <-p.rx: + if !ok { + return 0, false + } + n := copy(b, *buf) + p.pool.Put(buf) + return n, true + default: + return 0, false + } +} + func (p *packetPipe) readDeadline() (time.Time, <-chan struct{}) { p.deadline.mu.Lock() defer p.deadline.mu.Unlock()