Move frame wrapper

mtproto/wrappers/frame.go

@@ -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,
-	}
-}

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)
 

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,
+	}
+}