From 620561bcafdc504bf98ea6e8cbd543ed41571b91 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 25 Dec 2016 00:42:03 +0100
Subject: [PATCH] remove recover on stream write

---
 transport/ray/direct.go | 60 +++++++++++++++++++++++++++++------------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/transport/ray/direct.go b/transport/ray/direct.go
index f2d45b821..134bb44da 100644
--- a/transport/ray/direct.go
+++ b/transport/ray/direct.go
@@ -40,47 +40,73 @@ func (v *directRay) InboundOutput() InputStream {
 }
 
 type Stream struct {
-	buffer chan *buf.Buffer
+	buffer    chan *buf.Buffer
+	srcClose  chan bool
+	destClose chan bool
 }
 
 func NewStream() *Stream {
 	return &Stream{
-		buffer: make(chan *buf.Buffer, bufferSize),
+		buffer:    make(chan *buf.Buffer, bufferSize),
+		srcClose:  make(chan bool),
+		destClose: make(chan bool),
 	}
 }
 
 func (v *Stream) Read() (*buf.Buffer, error) {
-	buffer, open := <-v.buffer
-	if !open {
-		return nil, io.EOF
+	select {
+	case <-v.destClose:
+		return nil, io.ErrClosedPipe
+	case b := <-v.buffer:
+		return b, nil
+	default:
+		select {
+		case b := <-v.buffer:
+			return b, nil
+		case <-v.srcClose:
+			return nil, io.EOF
+		}
 	}
-	return buffer, nil
 }
 
 func (v *Stream) Write(data *buf.Buffer) (err error) {
-	defer func() {
-		if r := recover(); r != nil {
-			err = io.ErrClosedPipe
+	select {
+	case <-v.destClose:
+		return io.ErrClosedPipe
+	case <-v.srcClose:
+		return io.ErrClosedPipe
+	default:
+		select {
+		case <-v.destClose:
+			return io.ErrClosedPipe
+		case <-v.srcClose:
+			return io.ErrClosedPipe
+		case v.buffer <- data:
+			return nil
 		}
-	}()
-
-	v.buffer <- data
-	return nil
+	}
 }
 
 func (v *Stream) Close() {
 	defer swallowPanic()
 
-	close(v.buffer)
+	close(v.srcClose)
 }
 
 func (v *Stream) Release() {
 	defer swallowPanic()
 
-	close(v.buffer)
+	close(v.destClose)
+	v.Close()
 
-	for b := range v.buffer {
-		b.Release()
+	n := len(v.buffer)
+	for i := 0; i < n; i++ {
+		select {
+		case b := <-v.buffer:
+			b.Release()
+		default:
+			return
+		}
 	}
 }