diff options
| -rw-r--r-- | server/fsm.go | 10 | ||||
| -rw-r--r-- | server/fsm_test.go | 14 |
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 } |