Darien Raymond 7 yıl önce
ebeveyn
işleme
384844f898
1 değiştirilmiş dosya ile 8 ekleme ve 5 silme
  1. 8 5
      app/proxyman/inbound/worker.go

+ 8 - 5
app/proxyman/inbound/worker.go

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