From 1d93bef9130d81c6ea808ab3eb80aa04ff00f0db Mon Sep 17 00:00:00 2001 From: 9seconds Date: Sat, 7 Jul 2018 17:18:42 +0300 Subject: [PATCH] Move frame wrapper --- mtproto/wrappers/frame.go | 134 ------------------------------- wrappers/mtproto_abridged.go | 18 ++--- wrappers/mtproto_frame.go | 151 +++++++++++++++++++++++++++++++++++ 3 files changed, 160 insertions(+), 143 deletions(-) delete mode 100644 mtproto/wrappers/frame.go create mode 100644 wrappers/mtproto_frame.go diff --git a/mtproto/wrappers/frame.go b/mtproto/wrappers/frame.go deleted file mode 100644 index 633c4ad..0000000 --- a/mtproto/wrappers/frame.go +++ /dev/null @@ -1,134 +0,0 @@ -package wrappers - -import ( - "bytes" - "crypto/aes" - "encoding/binary" - "hash/crc32" - "io" - "io/ioutil" - "net" - - "github.com/juju/errors" - - "github.com/9seconds/mtg/wrappers" -) - -// Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] } -const ( - frameRWCMinMessageLength = 12 - frameRWCMaxMessageLength = 16777216 -) - -var frameRWCPadding = []byte{0x04, 0x00, 0x00, 0x00} - -type FrameRWC struct { - wrappers.BufferedReader - - conn wrappers.ReadWriteCloserWithAddr - readSeqNo int32 - writeSeqNo int32 -} - -func (f *FrameRWC) Write(buf []byte) (int, error) { - writeBuf := &bytes.Buffer{} - - // 4 - len bytes - // 4 - seq bytes - // . - message - // 4 - crc32 - messageLength := 4 + 4 + len(buf) + 4 - paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize - writeBuf.Grow(messageLength + paddingLength) - - binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength)) - binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo) - writeBuf.Write(buf) - f.writeSeqNo++ - - checksum := crc32.ChecksumIEEE(writeBuf.Bytes()) - binary.Write(writeBuf, binary.LittleEndian, checksum) - writeBuf.Write(bytes.Repeat(frameRWCPadding, paddingLength/4)) - - _, err := f.conn.Write(writeBuf.Bytes()) - return len(buf), err -} - -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() - 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) { - break - } - } - - messageLength := binary.LittleEndian.Uint32(buf.Bytes()) - if messageLength%4 != 0 || messageLength < frameRWCMinMessageLength || messageLength > frameRWCMaxMessageLength { - return errors.Errorf("Incorrect frame message length %d", messageLength) - } - - buf.Reset() - 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") - } - - var seqNo int32 - 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, _ := 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) - - } - f.Buffer.Write(data) - - return nil - }) -} - -func (f *FrameRWC) Close() error { - return f.conn.Close() -} - -func (f *FrameRWC) LocalAddr() *net.TCPAddr { - return f.conn.LocalAddr() -} - -func (f *FrameRWC) RemoteAddr() *net.TCPAddr { - return f.conn.RemoteAddr() -} - -func (f *FrameRWC) SocketID() string { - return f.conn.SocketID() -} - -func NewFrameRWC(conn wrappers.ReadWriteCloserWithAddr, seqNo int32) wrappers.ReadWriteCloserWithAddr { - return &FrameRWC{ - BufferedReader: wrappers.NewBufferedReader(), - conn: conn, - readSeqNo: seqNo, - writeSeqNo: seqNo, - } -} diff --git a/wrappers/mtproto_abridged.go b/wrappers/mtproto_abridged.go index b9950c3..1aa7f8c 100644 --- a/wrappers/mtproto_abridged.go +++ b/wrappers/mtproto_abridged.go @@ -12,9 +12,9 @@ import ( ) const ( - abridgedSmallPacketLength = 0x7f - abridgedQuickAckLength = 0x80 - abridgedLargePacketLength = 16777216 // 256 ^ 3 + mtprotoAbridgedSmallPacketLength = 0x7f + mtprotoAbridgedQuickAckLength = 0x80 + mtprotoAbridgedLargePacketLength = 16777216 // 256 ^ 3 ) type MTProtoAbridged struct { @@ -46,13 +46,13 @@ func (m *MTProtoAbridged) Read() ([]byte, error) { "counter", m.readCounter, ) - if msgLength >= abridgedQuickAckLength { + if msgLength >= mtprotoAbridgedQuickAckLength { m.opts.ReadHacks.QuickAck = true - msgLength -= abridgedQuickAckLength + msgLength -= mtprotoAbridgedQuickAckLength } msgLength32 := uint32(msgLength) - if msgLength == abridgedSmallPacketLength { + if msgLength == mtprotoAbridgedSmallPacketLength { if _, err := io.CopyN(buf, m.conn, 3); err != nil { return nil, errors.Annotate(err, "Cannot read the correct message length") } @@ -96,19 +96,19 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) { packetLength := len(p) / 4 switch { - case packetLength < abridgedSmallPacketLength: + case packetLength < mtprotoAbridgedSmallPacketLength: newData := append([]byte{byte(packetLength)}, p...) m.writeCounter++ return m.conn.Write(newData) - case packetLength < abridgedLargePacketLength: + case packetLength < mtprotoAbridgedLargePacketLength: length24 := utils.ToUint24(uint32(packetLength)) buf := &bytes.Buffer{} buf.Grow(1 + 3 + len(p)) - buf.WriteByte(byte(abridgedSmallPacketLength)) + buf.WriteByte(byte(mtprotoAbridgedSmallPacketLength)) buf.Write(length24[:]) buf.Write(p) diff --git a/wrappers/mtproto_frame.go b/wrappers/mtproto_frame.go new file mode 100644 index 0000000..d3eff6f --- /dev/null +++ b/wrappers/mtproto_frame.go @@ -0,0 +1,151 @@ +package wrappers + +import ( + "bytes" + "crypto/aes" + "encoding/binary" + "hash/crc32" + "io" + "io/ioutil" + "net" + + "github.com/juju/errors" +) + +const ( + mtprotoFrameMinMessageLength = 12 + mtprotoFrameMaxMessageLength = 16777216 +) + +var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00} + +type MTProtoFrame struct { + conn WrapStreamReadWriteCloser + readSeqNo int32 + writeSeqNo int32 +} + +func (m *MTProtoFrame) Read() ([]byte, error) { + buf := &bytes.Buffer{} + sum := crc32.NewIEEE() + writer := io.MultiWriter(buf, sum) + + for { + buf.Reset() + sum.Reset() + if _, err := io.CopyN(writer, m.conn, 4); err != nil { + return nil, errors.Annotate(err, "Cannot read frame padding") + } + if !bytes.Equal(buf.Bytes(), mtprotoFramePadding) { + break + } + } + + messageLength := binary.LittleEndian.Uint32(buf.Bytes()) + m.LogDebug("Read MTProto frame", + "messageLength", messageLength, + "sequence_number", m.readSeqNo, + ) + if messageLength%4 != 0 || messageLength < mtprotoFrameMinMessageLength || messageLength > mtprotoFrameMaxMessageLength { + return nil, errors.Errorf("Incorrect frame message length %d", messageLength) + } + + buf.Reset() + buf.Grow(int(messageLength) - 4 - 4) + if _, err := io.CopyN(writer, m.conn, int64(messageLength)-4-4); err != nil { + return nil, errors.Annotate(err, "Cannot read the message frame") + } + + var seqNo int32 + binary.Read(buf, binary.LittleEndian, &seqNo) + if seqNo != m.readSeqNo { + return nil, errors.Errorf("Unexpected sequence number %d (wait for %d)", seqNo, m.readSeqNo) + } + + 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, m.conn, 4); err != nil { + return nil, errors.Annotate(err, "Cannot read checksum") + } + + checksum := binary.LittleEndian.Uint32(buf.Bytes()) + if checksum != sum.Sum32() { + return nil, errors.Errorf("CRC32 checksum mismatch. Wait for %d, got %d", sum.Sum32(), checksum) + } + + m.LogDebug("Read MTProto frame", + "messageLength", messageLength, + "sequence_number", m.readSeqNo, + "dataLength", len(data), + "checksum", checksum, + ) + m.readSeqNo++ + + return data, nil +} + +func (m *MTProtoFrame) Write(p []byte) (int, error) { + messageLength := 4 + 4 + len(p) + 4 + paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize + + buf := &bytes.Buffer{} + buf.Grow(messageLength + paddingLength) + + binary.Write(buf, binary.LittleEndian, uint32(messageLength)) + binary.Write(buf, binary.LittleEndian, m.writeSeqNo) + buf.Write(p) + + checksum := crc32.ChecksumIEEE(buf.Bytes()) + binary.Write(buf, binary.LittleEndian, checksum) + buf.Write(bytes.Repeat(mtprotoFramePadding, paddingLength/4)) + + m.LogDebug("Write MTProto frame", + "length", len(p), + "sequence_number", m.writeSeqNo, + "crc32", checksum, + "frame_length", buf.Len(), + ) + m.writeSeqNo++ + + _, err := m.conn.Write(buf.Bytes()) + + return len(p), err +} + +func (m *MTProtoFrame) LogDebug(msg string, data ...interface{}) { + m.conn.LogDebug(msg, data...) +} + +func (m *MTProtoFrame) LogInfo(msg string, data ...interface{}) { + m.conn.LogInfo(msg, data...) +} + +func (m *MTProtoFrame) LogWarn(msg string, data ...interface{}) { + m.conn.LogWarn(msg, data...) +} + +func (m *MTProtoFrame) LogError(msg string, data ...interface{}) { + m.conn.LogError(msg, data...) +} + +func (m *MTProtoFrame) LocalAddr() *net.TCPAddr { + return m.conn.LocalAddr() +} + +func (m *MTProtoFrame) RemoteAddr() *net.TCPAddr { + return m.conn.RemoteAddr() +} + +func (m *MTProtoFrame) Close() error { + return m.conn.Close() +} + +func NewMTProtoFrame(conn WrapStreamReadWriteCloser, seqNo int32) WrapPacketReadWriteCloser { + return &MTProtoFrame{ + conn: conn, + readSeqNo: seqNo, + writeSeqNo: seqNo, + } +}