pull/861/head
Darien Raymond 2018-02-02 22:35:18 +01:00
parent cf832a4272
commit 384844f898
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 8 additions and 5 deletions

View File

@ -121,6 +121,7 @@ type udpConn struct {
output func([]byte) (int, error) output func([]byte) (int, error)
remote net.Addr remote net.Addr
local net.Addr local net.Addr
ctx context.Context
cancel context.CancelFunc cancel context.CancelFunc
} }
@ -129,13 +130,14 @@ func (c *udpConn) updateActivity() {
} }
func (c *udpConn) Read(buf []byte) (int, error) { func (c *udpConn) Read(buf []byte) (int, error) {
in, open := <-c.input select {
if !open { case in := <-c.input:
defer in.Release()
c.updateActivity()
return copy(buf, in.Bytes()), nil
case <-c.ctx.Done():
return 0, io.EOF return 0, io.EOF
} }
defer in.Release()
c.updateActivity()
return copy(buf, in.Bytes()), nil
} }
// Write implements io.Writer. // Write implements io.Writer.
@ -236,6 +238,7 @@ func (w *udpWorker) callback(b *buf.Buffer, source net.Destination, originalDest
go func() { go func() {
ctx := w.ctx ctx := w.ctx
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
conn.ctx = ctx
conn.cancel = cancel conn.cancel = cancel
if originalDest.IsValid() { if originalDest.IsValid() {
ctx = proxy.ContextWithOriginalTarget(ctx, originalDest) ctx = proxy.ContextWithOriginalTarget(ctx, originalDest)