From 1f51887dbe2c1c6766841c9a804bf3ed5ac97d0b Mon Sep 17 00:00:00 2001 From: Dimitri Sokolyuk Date: Thu, 28 Jul 2016 01:00:34 +0200 Subject: Frame/Deframe --- peer/messages.go | 35 ++++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/peer/messages.go b/peer/messages.go index a76248c..397ca04 100644 --- a/peer/messages.go +++ b/peer/messages.go @@ -1,9 +1,6 @@ package peer -import ( - "encoding/binary" - "io" -) +import "io" type Message int @@ -41,9 +38,7 @@ type Handshake struct { } func (h Handshake) Encode(w io.Writer) { - l := int8(len(h.Proto)) - binary.Write(w, binary.BigEndian, l) - w.Write([]byte(h.Proto)) + Frame(w, []byte(h.Proto)) w.Write(h.Flags[:]) w.Write(h.InfoHash[:]) w.Write(h.PeerID[:]) @@ -51,14 +46,7 @@ func (h Handshake) Encode(w io.Writer) { func DecodeHandshake(r io.Reader) (Handshake, bool) { var h Handshake - var l int8 - binary.Read(r, binary.BigEndian, &l) - if int(l) != len(Proto) { - return h, false - } - proto := make([]byte, int(l)) - r.Read(proto) - h.Proto = string(proto) + h.Proto = string(Deframe(r)) if h.Proto != Proto { return h, false } @@ -67,3 +55,20 @@ func DecodeHandshake(r io.Reader) (Handshake, bool) { r.Read(h.PeerID[:]) return h, true } + +func Frame(w io.Writer, b []byte) { + l := len(b) + w.Write([]byte{byte(l)}) + w.Write(b) +} + +func Deframe(r io.Reader) []byte { + l := make([]byte, 1) + r.Read(l) + if sz := int(l[0]); sz > 0 { + b := make([]byte, sz) + r.Read(b) + return b + } + return nil +} -- cgit v1.2.3