Small changes for abridged/intermediate

This commit is contained in:
9seconds
2018-07-05 16:37:25 +03:00
parent 4210cf014e
commit 154ff70bf7
3 changed files with 14 additions and 12 deletions
+7 -7
View File
@@ -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())
-1
View File
@@ -4,7 +4,6 @@ import (
"bytes"
"crypto/aes"
"encoding/binary"
"fmt"
"hash/crc32"
"io"
"io/ioutil"
+7 -4
View File
@@ -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")
}