Fix short read and report unexpected EOFs decoding peer protocol

[btrtrc.git] / peer_protocol / protocol.go

diff --git a/peer_protocol/protocol.go b/peer_protocol/protocol.go
index 281d84869a9666929ccedc6aa91081bc256dd5ae..a7fcbfc2b2da780121f61329eb476d9ea6bf5b3d 100644 (file)
--- a/peer_protocol/protocol.go
+++ b/peer_protocol/protocol.go
@@ -93,7 +93,7 @@ func (msg Message) MarshalBinary() (data []byte, err error) {
 
 type Decoder struct {
        R         *bufio.Reader
-       MaxLength Integer
+       MaxLength Integer // TODO: Should this include the length header or not?
 }
 
 func (d *Decoder) Decode(msg *Message) (err error) {
@@ -106,6 +106,14 @@ func (d *Decoder) Decode(msg *Message) (err error) {
                return errors.New("message too long")
        }
        r := bufio.NewReader(io.LimitReader(d.R, int64(length)))
+       defer func() {
+               written, _ := io.Copy(ioutil.Discard, r)
+               if written != 0 && err == nil {
+                       err = fmt.Errorf("short read on message type %d, left %d bytes", msg.Type, written)
+               } else if err == io.EOF {
+                       err = io.ErrUnexpectedEOF
+               }
+       }()
        if length == 0 {
                msg.Keepalive = true
                return
@@ -116,12 +124,6 @@ func (d *Decoder) Decode(msg *Message) (err error) {
                return
        }
        msg.Type = MessageType(c)
-       defer func() {
-               written, _ := io.Copy(ioutil.Discard, r)
-               if written != 0 && err != nil {
-                       err = fmt.Errorf("short read on message type %d, left %d bytes", msg.Type, written)
-               }
-       }()
        switch msg.Type {
        case Choke, Unchoke, Interested, NotInterested:
                return
@@ -152,9 +154,6 @@ func (d *Decoder) Decode(msg *Message) (err error) {
        default:
                err = fmt.Errorf("unknown message type %#v", c)
        }
-       if err != nil {
-               err = fmt.Errorf("decoding type %d: %s", msg.Type, err)
-       }
        return
 }