diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go index 9a77c73..784762f 100644 --- a/mtproto/wrappers/frame.go +++ b/mtproto/wrappers/frame.go @@ -6,6 +6,7 @@ import ( "encoding/binary" "hash/crc32" "io" + "io/ioutil" "net" "github.com/juju/errors" @@ -56,9 +57,13 @@ func (f *FrameRWC) Write(buf []byte) (int, error) { func (f *FrameRWC) Read(p []byte) (int, error) { return f.BufferedRead(p, func() error { buf := &bytes.Buffer{} + sum := crc32.NewIEEE() + writer := io.MultiWriter(buf, sum) + for { buf.Reset() - if _, err := io.CopyN(buf, f.conn, 4); err != nil { + sum.Reset() + if _, err := io.CopyN(writer, f.conn, 4); err != nil { return errors.Annotate(err, "Cannot read frame padding") } if !bytes.Equal(buf.Bytes(), frameRWCPadding[:]) { @@ -70,25 +75,29 @@ func (f *FrameRWC) Read(p []byte) (int, error) { if messageLength%4 != 0 || messageLength < frameRWCMinMessageLength || messageLength > frameRWCMaxMessageLength { return errors.Errorf("Incorrect frame message length %d", messageLength) } - sum := crc32.NewIEEE() - sum.Write(buf.Bytes()) buf.Reset() - buf.Grow(int(messageLength) - 4) // -4 because we already read the first number - if _, err := io.CopyN(buf, f.conn, int64(messageLength)-4); err != nil { + buf.Grow(int(messageLength) - 4 - 4) + if _, err := io.CopyN(writer, f.conn, int64(messageLength)-4-4); err != nil { return errors.Annotate(err, "Cannot read the message frame") } - sum.Write(buf.Bytes()) var seqNo int32 - binary.Read(buf, binary.LittleEndian, seqNo) + binary.Read(buf, binary.LittleEndian, &seqNo) if seqNo != f.readSeqNo { return errors.Errorf("Unexpected sequence number %d (wait for %d)", seqNo, f.readSeqNo) } f.readSeqNo++ - data := buf.Bytes()[:int(messageLength)-4-4-4] - checksum := binary.LittleEndian.Uint32(buf.Bytes()[int(messageLength)-4-4-4:]) + data, _ := ioutil.ReadAll(buf) + buf.Reset() + // write to buf, not to writer. This is because we are going to fetch + // crc32 checksum. + if _, err := io.CopyN(buf, f.conn, 4); err != nil { + return errors.Annotate(err, "Cannot read checksum") + } + checksum := binary.LittleEndian.Uint32(buf.Bytes()) + if checksum != sum.Sum32() { return errors.Errorf("CRC32 checksum mismatch. Wait for %d, got %d", sum.Sum32(), checksum) diff --git a/telegram/middle.go b/telegram/middle.go index 81975ca..a747b78 100644 --- a/telegram/middle.go +++ b/telegram/middle.go @@ -2,7 +2,6 @@ package telegram import ( "io" - "io/ioutil" "net" "net/http" "sync" @@ -84,11 +83,13 @@ func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.RPCNonceReque } func (t *middleTelegram) receiveRPCNonceResponse(conn io.Reader, req *rpc.RPCNonceRequest) (*rpc.RPCNonceResponse, error) { - ans, err := ioutil.ReadAll(conn) + var ans [128]byte + + n, err := conn.Read(ans[:]) if err != nil { return nil, errors.Annotate(err, "Cannot read RPC nonce response") } - rpcNonceResp, err := rpc.NewRPCNonceResponse(ans) + rpcNonceResp, err := rpc.NewRPCNonceResponse(ans[:n]) if err != nil { return nil, errors.Annotate(err, "Cannot initialize RPC nonce response") } @@ -109,11 +110,13 @@ func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.RPCHandsh } func (t *middleTelegram) receiveRPCHandshakeResponse(conn io.Reader, req *rpc.RPCHandshakeRequest) (*rpc.RPCHandshakeResponse, error) { - ans, err := ioutil.ReadAll(conn) + var ans [128]byte + + n, err := conn.Read(ans[:]) if err != nil { return nil, errors.Annotate(err, "Cannot read RPC handshake response") } - rpcHandshakeResp, err := rpc.NewRPCHandshakeResponse(ans) + rpcHandshakeResp, err := rpc.NewRPCHandshakeResponse(ans[:n]) if err != nil { return nil, errors.Annotate(err, "Cannot initialize RPC handshake response") }