fsm: fix homemade readAll()

Fix a received data corruption bug when a partial read happens. Let's
use io.ReadFull() rather than the homemade one.

Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 8 insertions, 16 deletions

diff --git a/server/fsm.go b/server/fsm.go
index 6b033060..42f2005a 100644
--- a/server/fsm.go
+++ b/server/fsm.go
@@ -21,6 +21,7 @@ import (
 	"github.com/osrg/gobgp/config"
 	"github.com/osrg/gobgp/packet"
 	"gopkg.in/tomb.v2"
+	"io"
 	"net"
 	"time"
 )
@@ -327,12 +328,9 @@ func buildopen(global *config.Global, peerConf *config.Neighbor) *bgp.BGPMessage
 
 func readAll(conn net.Conn, length int) ([]byte, error) {
 	buf := make([]byte, length)
-	for cur := 0; cur < length; {
-		if num, err := conn.Read(buf); err != nil {
-			return nil, err
-		} else {
-			cur += num
-		}
+	_, err := io.ReadFull(conn, buf)
+	if err != nil {
+		return nil, err
 	}
 	return buf, nil
 }
diff --git a/server/fsm_test.go b/server/fsm_test.go
index 7e321787..8c40d81c 100644
--- a/server/fsm_test.go
+++ b/server/fsm_test.go
@@ -65,25 +65,19 @@ func (m *MockConnection) Read(buf []byte) (int, error) {
 		m.currentCh = <-m.recvCh
 	}
 
-	rest := len(buf) - m.readBytes
+	length := 0
+	rest := len(buf)
 	for i := 0; i < rest; i++ {
 		if len(m.currentCh) > 0 {
 			val := <-m.currentCh
-			buf[m.readBytes] = val
-			m.readBytes += 1
+			buf[i] = val
+			length++
 		} else {
 			m.currentCh = nil
 			break
 		}
 	}
 
-	length := 0
-	if m.readBytes == len(buf) {
-		m.readBytes = 0
-		length = len(buf)
-	} else {
-		length = m.readBytes
-	}
 	fmt.Printf("%d bytes read from peer\n", length)
 	return length, nil
 }