diff --git a/mtproto/wrappers/abridged.go b/mtproto/wrappers/abridged.go index 66b440d..8d38a97 100644 --- a/mtproto/wrappers/abridged.go +++ b/mtproto/wrappers/abridged.go @@ -2,7 +2,6 @@ package wrappers import ( "bytes" - "encoding/binary" "io" "net" @@ -29,10 +28,14 @@ type AbridgedReadWriteCloserWithAddr struct { func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) { return a.BufferedRead(p, func() error { - var msgLength uint8 - if err := binary.Read(a.conn, binary.LittleEndian, &msgLength); err != nil { + buf := &bytes.Buffer{} + buf.Grow(3) + + if _, err := io.CopyN(buf, a.conn, 1); err != nil { return errors.Annotate(err, "Cannot read message length") } + msgLength := uint8(buf.Bytes()[0]) + buf.Reset() a.opts.QuickAck = false if msgLength >= abridgedQuickAckLength { @@ -42,11 +45,8 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) { msgLength32 := uint32(msgLength) if msgLength == abridgedSmallPacketLength { - buf := &bytes.Buffer{} - buf.Grow(3) - if _, err := io.CopyN(buf, a.conn, 3); err != nil { - return errors.Annotate(err, "Cannot read correct message length") + return errors.Annotate(err, "Cannot read the correct message length") } number := uint24{} copy(number[:], buf.Bytes()) diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go index 62406e7..907dbde 100644 --- a/mtproto/wrappers/frame.go +++ b/mtproto/wrappers/frame.go @@ -4,7 +4,6 @@ import ( "bytes" "crypto/aes" "encoding/binary" - "fmt" "hash/crc32" "io" "io/ioutil" diff --git a/mtproto/wrappers/intermediate.go b/mtproto/wrappers/intermediate.go index 1c2a6ab..4e35223 100644 --- a/mtproto/wrappers/intermediate.go +++ b/mtproto/wrappers/intermediate.go @@ -23,18 +23,21 @@ type IntermediateReadWriteCloserWithAddr struct { func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) { return i.BufferedRead(p, func() error { - var length uint32 - if err := binary.Read(i.conn, binary.LittleEndian, &length); err != nil { + buf := &bytes.Buffer{} + buf.Grow(4) + + if _, err := io.CopyN(buf, i.conn, 4); err != nil { return errors.Annotate(err, "Cannot read message length") } + length := binary.LittleEndian.Uint32(buf.Bytes()) + buf.Reset() + buf.Grow(int(length)) if length > intermediateQuickAckLength { i.opts.QuickAck = true length -= intermediateQuickAckLength } - buf := &bytes.Buffer{} - buf.Grow(int(length)) if _, err := io.CopyN(buf, i.conn, int64(length)); err != nil { return errors.Annotate(err, "Cannot read the message") }