From 1d05a7b82c4d84876cd1b7ce6fc4490c52c5d48a Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Fri, 2 Sep 2016 10:26:36 +0800
Subject: [PATCH] take back write lock

---
 session.go      | 24 +++++++++++++++---------
 session_test.go | 12 ++++++------
 stream.go       | 18 ++++++++----------
 3 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/session.go b/session.go
index 17cae14..1cf7836 100644
--- a/session.go
+++ b/session.go
@@ -21,7 +21,8 @@ const (
 
 // Session defines a multiplexed connection for streams
 type Session struct {
-	conn io.ReadWriteCloser
+	conn     io.ReadWriteCloser
+	sendLock sync.Mutex
 
 	config       *Config
 	nextStreamID uint32 // next stream identifier
@@ -76,7 +77,7 @@ func (s *Session) OpenStream() (*Stream, error) {
 	s.streams[sid] = stream
 	s.streamLock.Unlock()
 
-	s.sendFrame(newFrame(cmdSYN, sid))
+	s.writeFrame(newFrame(cmdSYN, sid))
 	return stream, nil
 }
 
@@ -105,7 +106,7 @@ func (s *Session) Close() error {
 			s.streams[k].Close()
 		}
 		s.streamLock.Unlock()
-		s.sendFrame(newFrame(cmdTerminate, 0))
+		s.writeFrame(newFrame(cmdTerminate, 0))
 		s.conn.Close()
 		close(s.die)
 	}
@@ -203,6 +204,8 @@ func (s *Session) recvLoop() {
 		case <-s.tbf:
 			s.tbf <- struct{}{}
 			if f, err := s.readFrame(buffer); err == nil {
+				atomic.StoreInt32(&s.dataReady, 1)
+
 				switch f.cmd {
 				case cmdNOP:
 				case cmdTerminate:
@@ -214,7 +217,7 @@ func (s *Session) recvLoop() {
 						s.streams[f.sid] = newStream(f.sid, s.config.MaxFrameSize, s)
 						s.chAccepts <- s.streams[f.sid]
 					} else { // stream exists, RST the peer
-						s.sendFrame(newFrame(cmdRST, f.sid))
+						s.writeFrame(newFrame(cmdRST, f.sid))
 					}
 					s.streamLock.Unlock()
 				case cmdRST:
@@ -231,14 +234,13 @@ func (s *Session) recvLoop() {
 						stream.notifyReadEvent()
 						<-s.tbf // remove a token
 					} else { // stream is absent
-						s.sendFrame(newFrame(cmdRST, f.sid))
+						s.writeFrame(newFrame(cmdRST, f.sid))
 					}
 					s.streamLock.Unlock()
 				default:
 					s.Close()
 					return
 				}
-				atomic.StoreInt32(&s.dataReady, 1)
 			} else {
 				s.Close()
 				return
@@ -257,7 +259,7 @@ func (s *Session) keepalive() {
 	for {
 		select {
 		case <-tickerPing.C:
-			s.sendFrame(newFrame(cmdNOP, 0))
+			s.writeFrame(newFrame(cmdNOP, 0))
 		case <-tickerTimeout.C:
 			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) && len(s.tbf) == s.config.MaxFrameTokens {
 				s.Close()
@@ -269,7 +271,11 @@ func (s *Session) keepalive() {
 	}
 }
 
-func (s *Session) sendFrame(f Frame) {
+// writeFrame writes the frame to the underlying connection, and returns len(f.data) if successful
+func (s *Session) writeFrame(f Frame) (n int, err error) {
 	bts, _ := f.MarshalBinary()
-	s.conn.Write(bts)
+	s.sendLock.Lock()
+	_, err = s.conn.Write(bts)
+	s.sendLock.Unlock()
+	return len(f.data), err
 }
diff --git a/session_test.go b/session_test.go
index 532da67..3b0b5d7 100644
--- a/session_test.go
+++ b/session_test.go
@@ -224,8 +224,8 @@ func TestKeepAliveTimeout(t *testing.T) {
 	}
 
 	config := DefaultConfig()
-	config.KeepAliveInterval = 1
-	config.KeepAliveTimeout = 2
+	config.KeepAliveInterval = time.Second
+	config.KeepAliveTimeout = 2 * time.Second
 	session, _ := Client(cli, config)
 	<-time.After(3 * time.Second)
 	if session.IsClosed() != true {
@@ -339,7 +339,7 @@ func TestRandomFrame(t *testing.T) {
 	session, _ = Client(cli, nil)
 	for i := 0; i < 100; i++ {
 		f := newFrame(cmdSYN, 1000)
-		session.sendFrame(f)
+		session.writeFrame(f)
 	}
 	cli.Close()
 
@@ -352,7 +352,7 @@ func TestRandomFrame(t *testing.T) {
 	session, _ = Client(cli, nil)
 	for i := 0; i < 100; i++ {
 		f := newFrame(allcmds[rand.Int()%len(allcmds)], rand.Uint32())
-		session.sendFrame(f)
+		session.writeFrame(f)
 	}
 	cli.Close()
 
@@ -364,7 +364,7 @@ func TestRandomFrame(t *testing.T) {
 	session, _ = Client(cli, nil)
 	for i := 0; i < 100; i++ {
 		f := newFrame(byte(rand.Uint32()), rand.Uint32())
-		session.sendFrame(f)
+		session.writeFrame(f)
 	}
 	cli.Close()
 
@@ -377,7 +377,7 @@ func TestRandomFrame(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		f := newFrame(byte(rand.Uint32()), rand.Uint32())
 		f.ver = byte(rand.Uint32())
-		session.sendFrame(f)
+		session.writeFrame(f)
 	}
 	cli.Close()
 
diff --git a/stream.go b/stream.go
index b267a6b..188e75f 100644
--- a/stream.go
+++ b/stream.go
@@ -1,7 +1,6 @@
 package smux
 
 import (
-	"bytes"
 	"sync"
 
 	"github.com/pkg/errors"
@@ -75,16 +74,15 @@ func (s *Stream) Write(b []byte) (n int, err error) {
 	}
 
 	frames := s.split(b, cmdPSH, s.id)
-	var combined bytes.Buffer
+	sum := 0
 	for k := range frames {
-		bts, _ := frames[k].MarshalBinary()
-		combined.Write(bts)
+		fn, ferr := s.sess.writeFrame(frames[k])
+		sum += fn
+		if ferr != nil {
+			return sum, ferr
+		}
 	}
-
-	if _, err = s.sess.conn.Write(combined.Bytes()); err != nil {
-		return 0, err
-	}
-	return len(b), nil
+	return sum, nil
 }
 
 // Close implements io.ReadWriteCloser
@@ -98,7 +96,7 @@ func (s *Stream) Close() error {
 	default:
 		close(s.die)
 		s.sess.streamClosed(s.id)
-		s.sess.sendFrame(newFrame(cmdRST, s.id))
+		s.sess.writeFrame(newFrame(cmdRST, s.id))
 	}
 	return nil
 }