From 6f8177ccd6fd6c88bc8016068cb724deb90ea4cf Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Thu, 1 Sep 2016 14:30:07 +0800
Subject: [PATCH] optimize recvLoop

---
 session.go |  9 ++++++++-
 stream.go  | 21 +++++++--------------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/session.go b/session.go
index 8e8f3f0..8952c48 100644
--- a/session.go
+++ b/session.go
@@ -215,7 +215,12 @@ func (s *Session) recvLoop() {
 						s.sendFrame(newFrame(cmdRST, f.sid))
 					}
 					s.tbf <- struct{}{}
-				default:
+				case cmdRST:
+					if _, ok := s.streams[f.sid]; ok {
+						s.streams[f.sid].Close()
+					}
+					s.tbf <- struct{}{}
+				case cmdPSH:
 					if _, ok := s.streams[f.sid]; ok {
 						s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f)
 						select {
@@ -226,6 +231,8 @@ func (s *Session) recvLoop() {
 						s.sendFrame(newFrame(cmdRST, f.sid))
 						s.tbf <- struct{}{}
 					}
+				default:
+					s.sendFrame(newFrame(cmdRST, f.sid))
 				}
 				s.mu.Unlock()
 				atomic.StoreInt32(&s.dataReady, 1)
diff --git a/stream.go b/stream.go
index cf297fd..4d51734 100644
--- a/stream.go
+++ b/stream.go
@@ -2,7 +2,6 @@ package smux
 
 import (
 	"bytes"
-	"io"
 	"sync"
 
 	"github.com/pkg/errors"
@@ -47,20 +46,14 @@ READ:
 		return n, nil
 	}
 
-	if f := s.sess.nioread(s.id); f != nil {
-		switch f.cmd {
-		case cmdPSH:
-			n = copy(b, f.data)
-			if len(f.data) > n {
-				s.buffer = make([]byte, len(f.data)-n)
-				copy(s.buffer, f.data[n:])
-			}
-			s.rlock.Unlock()
-			return n, nil
-		default:
-			s.rlock.Unlock()
-			return 0, io.EOF
+	if f := s.sess.nioread(s.id); f != nil && f.cmd == cmdPSH {
+		n = copy(b, f.data)
+		if len(f.data) > n {
+			s.buffer = make([]byte, len(f.data)-n)
+			copy(s.buffer, f.data[n:])
 		}
+		s.rlock.Unlock()
+		return n, nil
 	}
 
 	s.rlock.Unlock()