From 5d4c4e586972189cede7c4b97350824550fa5a67 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Sat, 3 Sep 2016 15:38:31 +0800
Subject: [PATCH] a better fix for rst

---
 session.go | 41 +++++++++++++++++++++--------------------
 stream.go  | 35 ++++++++++++-----------------------
 2 files changed, 33 insertions(+), 43 deletions(-)

diff --git a/session.go b/session.go
index 0659772..96e1684 100644
--- a/session.go
+++ b/session.go
@@ -30,9 +30,9 @@ type Session struct {
 	bucket     int32
 	bucketCond *sync.Cond
 
-	frameQueues map[uint32][]Frame // stream input frame queue
-	streams     map[uint32]*Stream // all streams in this session
-	streamLock  sync.Mutex         // locks streams && frameQueues
+	streamBuffers map[uint32][]byte  // stream input buffer
+	streams       map[uint32]*Stream // all streams in this session
+	streamLock    sync.Mutex         // locks streams && frameQueues
 
 	die            chan struct{} // flag session has died
 	dieLock        sync.Mutex
@@ -48,7 +48,7 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
 	s.conn = conn
 	s.config = config
 	s.streams = make(map[uint32]*Stream)
-	s.frameQueues = make(map[uint32][]Frame)
+	s.streamBuffers = make(map[uint32][]byte)
 	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
 	s.chClosedStream = make(chan uint32, defaultCloseWait)
 	s.bucket = int32(config.MaxReceiveBuffer)
@@ -145,19 +145,16 @@ func (s *Session) streamClosed(sid uint32) {
 }
 
 // nonblocking read from session pool, for streams
-func (s *Session) nioread(sid uint32) *Frame {
+func (s *Session) nioread(sid uint32, p []byte) (n int) {
 	s.streamLock.Lock()
-	frames := s.frameQueues[sid]
-	if len(frames) > 0 {
-		f := frames[0]
-		s.frameQueues[sid] = frames[1:]
-		atomic.AddInt32(&s.bucket, int32(len(f.data)))
-		s.streamLock.Unlock()
+	n = copy(p, s.streamBuffers[sid])
+	if n > 0 {
+		s.streamBuffers[sid] = s.streamBuffers[sid][n:]
+		atomic.AddInt32(&s.bucket, int32(n))
 		s.bucketCond.Signal()
-		return &f
 	}
 	s.streamLock.Unlock()
-	return nil
+	return
 }
 
 // session read a frame from underlying connection
@@ -189,12 +186,11 @@ func (s *Session) monitor() {
 		case sid := <-s.chClosedStream:
 			s.streamLock.Lock()
 			delete(s.streams, sid)
-			fq := s.frameQueues[sid]
-			for k := range fq { // return remaining tokens to the bucket
-				atomic.AddInt32(&s.bucket, int32(len(fq[k].data)))
+			if n := len(s.streamBuffers[sid]); n > 0 { // return remaining tokens to the bucket
+				atomic.AddInt32(&s.bucket, int32(n))
+				s.bucketCond.Signal()
 			}
-			s.bucketCond.Signal()
-			delete(s.frameQueues, sid)
+			delete(s.streamBuffers, sid)
 			s.streamLock.Unlock()
 		case <-s.die:
 			return
@@ -235,12 +231,17 @@ func (s *Session) recvLoop() {
 				}
 				s.streamLock.Unlock()
 			case cmdRST:
-				fallthrough
+				s.streamLock.Lock()
+				if stream, ok := s.streams[f.sid]; ok {
+					stream.markRST()
+					stream.notifyReadEvent()
+				}
+				s.streamLock.Unlock()
 			case cmdPSH:
 				s.streamLock.Lock()
 				if stream, ok := s.streams[f.sid]; ok {
 					atomic.AddInt32(&s.bucket, -int32(len(f.data)))
-					s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f)
+					s.streamBuffers[f.sid] = append(s.streamBuffers[f.sid], f.data...)
 					stream.notifyReadEvent()
 				} else { // stream is absent
 					go s.writeFrame(newFrame(cmdRST, f.sid))
diff --git a/stream.go b/stream.go
index 1c1634b..7a07a5e 100644
--- a/stream.go
+++ b/stream.go
@@ -3,6 +3,7 @@ package smux
 import (
 	"bytes"
 	"sync"
+	"sync/atomic"
 
 	"github.com/pkg/errors"
 )
@@ -10,10 +11,9 @@ import (
 // Stream implements io.ReadWriteCloser
 type Stream struct {
 	id          uint32
+	rstflag     int32
 	sess        *Session
 	frameSize   int
-	rlock       sync.Mutex    // read lock
-	buffer      []byte        // temporary store of remaining frame.data
 	chReadEvent chan struct{} // notify a read event
 	die         chan struct{} // flag the stream has closed
 	dieLock     sync.Mutex
@@ -39,29 +39,13 @@ READ:
 	default:
 	}
 
-	s.rlock.Lock()
-	if len(s.buffer) > 0 {
-		n = copy(b, s.buffer)
-		s.buffer = s.buffer[n:]
-		s.rlock.Unlock()
-		return n, nil
+	if n = s.sess.nioread(s.id, b); n > 0 {
+		return n, err
+	} else if atomic.LoadInt32(&s.rstflag) == 1 {
+		s.Close()
+		return 0, errors.New(errBrokenPipe)
 	}
 
-	if f := s.sess.nioread(s.id); f != nil {
-		switch f.cmd {
-		case cmdPSH:
-			n = copy(b, f.data)
-			s.buffer = f.data[n:]
-			s.rlock.Unlock()
-			return
-		case cmdRST:
-			s.Close()
-			s.rlock.Unlock()
-			return 0, errors.New(errBrokenPipe)
-		}
-	}
-
-	s.rlock.Unlock()
 	select {
 	case <-s.chReadEvent:
 		goto READ
@@ -134,3 +118,8 @@ func (s *Stream) notifyReadEvent() {
 	default:
 	}
 }
+
+// mark this stream has benn reset
+func (s *Stream) markRST() {
+	atomic.StoreInt32(&s.rstflag, 1)
+}