1 files changed, 15 insertions, 54 deletions

diff --git a/peer.go b/peer.go
index a855fe2b9..d47af73de 100644
--- a/peer.go
+++ b/peer.go
@@ -2,57 +2,24 @@ package main
 
 import (
   "net"
-  "errors"
   "log"
+  "github.com/ethereum/ethwire-go"
 )
 
-type InMsg struct {
-  msgType   string  // Specifies how the encoded data should be interpreted
-  data      []byte  // RLP encoded data
-}
-
-type OutMsg struct {
-  msgType   string
-  data      []byte
-}
-
-func ReadMessage(conn net.Conn) (*InMsg, error) {
-  buff := make([]byte, 4069)
-
-  // Wait for a message from this peer
-  n, err := conn.Read(buff)
-  if err != nil {
-    return nil, err
-  } else if n == 0 {
-    return nil, errors.New("Empty message received")
-  }
-
-  // Read the header (MAX n)
-  // XXX The data specification is made up. This will change once more details have been released on the specification of the format
-  decoder := NewRlpDecoder(buff[:n])
-  t := decoder.Get(0).AsString()
-  // If the msgdata contains no data we throw an error and disconnect the peer
-  if t == "" {
-    return nil, errors.New("Data contained no data type")
-  }
-
-  return &InMsg{msgType: t, data: decoder.Get(1).AsBytes()}, nil
-}
-
 type Peer struct {
   // Server interface
   server      *Server
   // Net connection
   conn        net.Conn
   // Output queue which is used to communicate and handle messages
-  outputQueue chan OutMsg
+  outputQueue chan ethwire.InOutMsg
   // Quit channel
   quit        chan bool
 }
 
 func NewPeer(conn net.Conn, server *Server) *Peer {
   return &Peer{
-    outputQueue:       make(chan OutMsg, 1),  // Buffered chan of 1 is enough
+    outputQueue:       make(chan ethwire.InOutMsg, 1),  // Buffered chan of 1 is enough
     quit:              make(chan bool),
 
     server:            server,
@@ -62,7 +29,7 @@ func NewPeer(conn net.Conn, server *Server) *Peer {
 
 // Outputs any RLP encoded data to the peer
 func (p *Peer) QueueMessage(msgType string, data []byte) {
-  p.outputQueue <- OutMsg{msgType: msgType, data: data}
+  p.outputQueue <- ethwire.InOutMsg{MsgType: msgType, Data: data}
 }
 
 // Outbound message handler. Outbound messages are handled here
@@ -73,7 +40,13 @@ out:
     // Main message queue. All outbound messages are processed through here
     case msg := <-p.outputQueue:
       // TODO Message checking and handle accordingly
-      p.WriteMessage(msg)
+      err := ethwire.WriteMessage(p.conn, msg)
+      if err != nil {
+        log.Println(err)
+
+        // Stop the client if there was an error writing to it
+        p.Stop()
+      }
 
     // Break out of the for loop if a quit message is posted
     case <- p.quit:
@@ -82,18 +55,6 @@ out:
   }
 }
 
-// Write a message to the peer.
-func (p *Peer) WriteMessage(msg OutMsg) {
-  // Encode the type and the (RLP encoded) data for sending over the wire
-  encoded := Encode([]interface{}{ msg.msgType, msg.data })
-  // Write to the connection
-  _, err := p.conn.Write(encoded)
-  if err != nil {
-    log.Println(err)
-    p.Stop()
-  }
-}
-
 // Inbound handler. Inbound messages are received here and passed to the appropriate methods
 func (p *Peer) HandleInbound() {
   defer p.Stop()
@@ -101,7 +62,7 @@ func (p *Peer) HandleInbound() {
 out:
   for {
     // Wait for a message from the peer
-    msg, err := ReadMessage(p.conn)
+    msg, err := ethwire.ReadMessage(p.conn)
     if err != nil {
       log.Println(err)
 
@@ -109,8 +70,8 @@ out:
     }
 
     // TODO
-    data, _ := Decode(msg.data, 0)
-    log.Printf("%s, %s\n", msg.msgType, data)
+    data, _ := Decode(msg.Data, 0)
+    log.Printf("%s, %s\n", msg.MsgType, data)
   }
 
   // Notify the out handler we're quiting
@@ -125,7 +86,7 @@ func (p *Peer) Start() {
 }
 
 func (p *Peer) Stop() {
-  defer p.conn.Close()
+  p.conn.Close()
 
   p.quit <- true
 }