From 421fe5bf76078e8fc666dce926e3c10baa44dbd6 Mon Sep 17 00:00:00 2001
From: Brian Tiger Chow <brian.holderchow@gmail.com>
Date: Thu, 11 Sep 2014 03:09:36 -0700
Subject: [PATCH] refactor(bitswap:message) add wrapper for proto

---
 bitswap/bitswap.go      | 21 +++++-------
 bitswap/message.go      | 38 +++++++++++++++++++++
 bitswap/message_test.go | 75 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 122 insertions(+), 12 deletions(-)
 create mode 100644 bitswap/message.go
 create mode 100644 bitswap/message_test.go

diff --git a/bitswap/bitswap.go b/bitswap/bitswap.go
index 22f97514cf8..05be4ce0a5f 100644
--- a/bitswap/bitswap.go
+++ b/bitswap/bitswap.go
@@ -115,13 +115,12 @@ func (bs *BitSwap) GetBlock(k u.Key, timeout time.Duration) (
 func (bs *BitSwap) getBlock(k u.Key, p *peer.Peer, timeout time.Duration) ([]byte, error) {
 	u.DOut("[%s] getBlock '%s' from [%s]\n", bs.peer.ID.Pretty(), k.Pretty(), p.ID.Pretty())
 
-	pmes := new(PBMessage)
-	pmes.Wantlist = []string{string(k)}
+	message := newMessage()
+	message.AppendWanted(k)
 
 	after := time.After(timeout)
 	resp := bs.listener.Listen(string(k), 1, timeout)
-	smes := swarm.NewMessage(p, pmes)
-	bs.meschan.Outgoing <- smes
+	bs.meschan.Outgoing <- message.ToSwarm(p)
 
 	select {
 	case resp_mes := <-resp:
@@ -149,11 +148,9 @@ func (bs *BitSwap) HaveBlock(blk *blocks.Block) error {
 }
 
 func (bs *BitSwap) SendBlock(p *peer.Peer, b *blocks.Block) {
-	pmes := new(PBMessage)
-	pmes.Blocks = [][]byte{b.Data}
-
-	swarm_mes := swarm.NewMessage(p, pmes)
-	bs.meschan.Outgoing <- swarm_mes
+	message := newMessage()
+	message.AppendBlock(b)
+	bs.meschan.Outgoing <- message.ToSwarm(p)
 }
 
 func (bs *BitSwap) handleMessages() {
@@ -257,14 +254,14 @@ func (bs *BitSwap) GetLedger(p *peer.Peer) *Ledger {
 }
 
 func (bs *BitSwap) SendWantList(wl KeySet) error {
-	pmes := new(PBMessage)
+	message := newMessage()
 	for k, _ := range wl {
-		pmes.Wantlist = append(pmes.Wantlist, string(k))
+		message.AppendWanted(k)
 	}
 
 	// Lets just ping everybody all at once
 	for _, ledger := range bs.partners {
-		bs.meschan.Outgoing <- swarm.NewMessage(ledger.Partner, pmes)
+		bs.meschan.Outgoing <- message.ToSwarm(ledger.Partner)
 	}
 
 	return nil
diff --git a/bitswap/message.go b/bitswap/message.go
new file mode 100644
index 00000000000..94bb82ef89b
--- /dev/null
+++ b/bitswap/message.go
@@ -0,0 +1,38 @@
+package bitswap
+
+import (
+	blocks "github.com/jbenet/go-ipfs/blocks"
+	peer "github.com/jbenet/go-ipfs/peer"
+	swarm "github.com/jbenet/go-ipfs/swarm"
+	u "github.com/jbenet/go-ipfs/util"
+)
+
+// message wraps a proto message for convenience
+type message struct {
+	pb PBMessage
+}
+
+func newMessageFromProto(pb PBMessage) *message {
+	return &message{pb: pb}
+}
+
+func newMessage() *message {
+	return new(message)
+}
+
+func (m *message) AppendWanted(k u.Key) {
+	m.pb.Wantlist = append(m.pb.Wantlist, string(k))
+}
+
+func (m *message) AppendBlock(b *blocks.Block) {
+	m.pb.Blocks = append(m.pb.Blocks, b.Data)
+}
+
+func (m *message) ToProto() *PBMessage {
+	cp := m.pb
+	return &cp
+}
+
+func (m *message) ToSwarm(p *peer.Peer) *swarm.Message {
+	return swarm.NewMessage(p, m.ToProto())
+}
diff --git a/bitswap/message_test.go b/bitswap/message_test.go
new file mode 100644
index 00000000000..bc52b5aa9ed
--- /dev/null
+++ b/bitswap/message_test.go
@@ -0,0 +1,75 @@
+package bitswap
+
+import (
+	"bytes"
+	"testing"
+
+	blocks "github.com/jbenet/go-ipfs/blocks"
+	u "github.com/jbenet/go-ipfs/util"
+)
+
+func TestAppendWanted(t *testing.T) {
+	const str = "foo"
+	m := newMessage()
+	m.AppendWanted(u.Key(str))
+
+	if !contains(m.ToProto().GetWantlist(), str) {
+		t.Fail()
+	}
+}
+
+func TestNewMessageFromProto(t *testing.T) {
+	const str = "a_key"
+	protoMessage := new(PBMessage)
+	protoMessage.Wantlist = []string{string(str)}
+	if !contains(protoMessage.Wantlist, str) {
+		t.Fail()
+	}
+	m := newMessageFromProto(*protoMessage)
+	if !contains(m.ToProto().GetWantlist(), str) {
+		t.Fail()
+	}
+}
+
+func TestAppendBlock(t *testing.T) {
+
+	strs := make([]string, 2)
+	strs = append(strs, "Celeritas")
+	strs = append(strs, "Incendia")
+
+	m := newMessage()
+	for _, str := range strs {
+		block, err := blocks.NewBlock([]byte(str))
+		if err != nil {
+			t.Fail()
+		}
+		m.AppendBlock(block)
+	}
+
+	// assert strings are in proto message
+	for _, blockbytes := range m.ToProto().GetBlocks() {
+		s := bytes.NewBuffer(blockbytes).String()
+		if !contains(strs, s) {
+			t.Fail()
+		}
+	}
+}
+
+func TestCopyProtoByValue(t *testing.T) {
+	const str = "foo"
+	m := newMessage()
+	protoBeforeAppend := m.ToProto()
+	m.AppendWanted(u.Key(str))
+	if contains(protoBeforeAppend.GetWantlist(), str) {
+		t.Fail()
+	}
+}
+
+func contains(s []string, x string) bool {
+	for _, a := range s {
+		if a == x {
+			return true
+		}
+	}
+	return false
+}