From 591b0f431df047b4c2ff4eadde1a65bce859d5a4 Mon Sep 17 00:00:00 2001
From: Dave Collins <davec@conformal.com>
Date: Mon, 3 Feb 2014 23:47:57 -0600
Subject: [PATCH] Switch over to new btcwire Read/WriteMessageN.

This commit adds byte counters to each peer using the new btcwire
ReadMessageN and WriteMessageN functions to obtain the number of bytes
read and written, respectively.  It also returns those byte counters via
the PeerInfo struct which is used to populate the RPC getpeerinfo reply.

Closes #83.
---
 peer.go   | 16 +++++++++++-----
 server.go |  8 ++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/peer.go b/peer.go
index f7e3bf05..b4c105da 100644
--- a/peer.go
+++ b/peer.go
@@ -137,6 +137,8 @@ type peer struct {
 	timeConnected      time.Time
 	lastSend           time.Time
 	lastRecv           time.Time
+	bytesRead          uint64
+	bytesWritten       uint64
 	inbound            bool
 	connected          int32
 	disconnect         int32 // only to be used atomically
@@ -951,11 +953,13 @@ func (p *peer) handlePongMsg(msg *btcwire.MsgPong) {
 }
 
 // readMessage reads the next bitcoin message from the peer with logging.
-func (p *peer) readMessage() (msg btcwire.Message, buf []byte, err error) {
-	msg, buf, err = btcwire.ReadMessage(p.conn, p.protocolVersion, p.btcnet)
+func (p *peer) readMessage() (btcwire.Message, []byte, error) {
+	n, msg, buf, err := btcwire.ReadMessageN(p.conn, p.protocolVersion, p.btcnet)
 	if err != nil {
-		return
+		p.bytesRead += uint64(n)
+		return nil, nil, err
 	}
+	p.bytesRead += uint64(n)
 
 	// Use closures to log expensive operations so they are only run when
 	// the logging level requires it.
@@ -975,7 +979,7 @@ func (p *peer) readMessage() (msg btcwire.Message, buf []byte, err error) {
 		return spew.Sdump(buf)
 	}))
 
-	return
+	return msg, buf, nil
 }
 
 // writeMessage sends a bitcoin Message to the peer with logging.
@@ -1019,12 +1023,14 @@ func (p *peer) writeMessage(msg btcwire.Message) {
 	}))
 
 	// Write the message to the peer.
-	err := btcwire.WriteMessage(p.conn, msg, p.protocolVersion, p.btcnet)
+	n, err := btcwire.WriteMessageN(p.conn, msg, p.protocolVersion, p.btcnet)
 	if err != nil {
+		p.bytesWritten += uint64(n)
 		p.Disconnect()
 		p.logError("Can't send message: %v", err)
 		return
 	}
+	p.bytesWritten += uint64(n)
 }
 
 // isAllowedByRegression returns whether or not the passed error is allowed by
diff --git a/server.go b/server.go
index 0f132e1e..b51e0413 100644
--- a/server.go
+++ b/server.go
@@ -264,8 +264,8 @@ type PeerInfo struct {
 	Services       string `json:"services"`
 	LastSend       int64  `json:"lastsend"`
 	LastRecv       int64  `json:"lastrecv"`
-	BytesSent      int    `json:"bytessent"`
-	BytesRecv      int    `json:"bytesrecv"`
+	BytesSent      uint64 `json:"bytessent"`
+	BytesRecv      uint64 `json:"bytesrecv"`
 	PingTime       int64  `json:"pingtime"`
 	PingWait       int64  `json:"pingwait,omitempty"`
 	ConnTime       int64  `json:"conntime"`
@@ -328,8 +328,8 @@ func (s *server) handleQuery(querymsg interface{}, state *peerState) {
 				Services:       fmt.Sprintf("%08d", p.services),
 				LastSend:       p.lastSend.Unix(),
 				LastRecv:       p.lastRecv.Unix(),
-				BytesSent:      0, // TODO(oga) we need this from wire.
-				BytesRecv:      0, // TODO(oga) we need this from wire.
+				BytesSent:      p.bytesWritten,
+				BytesRecv:      p.bytesRead,
 				ConnTime:       p.timeConnected.Unix(),
 				Version:        p.protocolVersion,
 				SubVer:         p.userAgent,