Bugfix framerwc

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)
 

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")
 	}