Add base wrappers for mtproto

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/cespare/xxhash v1.1.0
 	github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d
 	github.com/kr/pretty v0.1.0 // indirect
-	github.com/pkg/errors v0.8.1 // indirect
+	github.com/pkg/errors v0.8.1
 	github.com/prometheus/client_golang v1.1.0
 	github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect

mtproto/consts.go

@@ -0,0 +1,32 @@
+package mtproto
+
+// SeqNo* is the number of the sequence which have special meaning for
+// the Telegram.
+const (
+	SeqNoNonce     = -2
+	SeqNoHandshake = -1
+)
+
+// Different constants for RPC protocol
+var (
+	TagCloseExt     = []byte{0xa2, 0x34, 0xb6, 0x5e}
+	TagProxyAns     = []byte{0x0d, 0xda, 0x03, 0x44}
+	TagSimpleAck    = []byte{0x9b, 0x40, 0xac, 0x3b}
+	TagHandshake    = []byte{0xf5, 0xee, 0x82, 0x76}
+	TagNonce        = []byte{0xaa, 0x87, 0xcb, 0x7a}
+	TagProxyRequest = []byte{0xee, 0xf1, 0xce, 0x36}
+
+	NonceCryptoAES = []byte{0x01, 0x00, 0x00, 0x00}
+
+	HandshakeFlags = []byte{0x00, 0x00, 0x00, 0x00}
+
+	ProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
+	ProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
+
+	HandshakeSenderPID = []byte("IPIPPRPDTIME")
+	HandshakePeerPID   = []byte("IPIPPRPDTIME")
+
+	HandshakeRequest = append(TagHandshake,
+		append(HandshakeFlags,
+			append(HandshakeSenderPID, HandshakePeerPID...)...)...)
+)

mtproto/proxy_flags.go

@@ -0,0 +1,59 @@
+package mtproto
+
+import (
+	"encoding/binary"
+	"strings"
+)
+
+type proxyRequestFlags uint32
+
+const (
+	proxyRequestFlagsHasAdTag     proxyRequestFlags = 0x8
+	proxyRequestFlagsEncrypted    proxyRequestFlags = 0x2
+	proxyRequestFlagsMagic        proxyRequestFlags = 0x1000
+	proxyRequestFlagsExtMode2     proxyRequestFlags = 0x20000
+	proxyRequestFlagsIntermediate proxyRequestFlags = 0x20000000
+	proxyRequestFlagsAbdridged    proxyRequestFlags = 0x40000000
+	proxyRequestFlagsQuickAck     proxyRequestFlags = 0x80000000
+	proxyRequestFlagsPad          proxyRequestFlags = 0x8000000
+)
+
+var proxyRequestFlagsEncryptedPrefix [8]byte
+
+func (r proxyRequestFlags) Bytes() []byte {
+	converted := make([]byte, 4)
+	binary.LittleEndian.PutUint32(converted, uint32(r))
+
+	return converted
+}
+
+func (r proxyRequestFlags) String() string {
+	flags := make([]string, 0, 7)
+
+	if r&proxyRequestFlagsHasAdTag != 0 {
+		flags = append(flags, "HAS_AD_TAG")
+	}
+	if r&proxyRequestFlagsEncrypted != 0 {
+		flags = append(flags, "ENCRYPTED")
+	}
+	if r&proxyRequestFlagsMagic != 0 {
+		flags = append(flags, "MAGIC")
+	}
+	if r&proxyRequestFlagsExtMode2 != 0 {
+		flags = append(flags, "EXT_MODE_2")
+	}
+	if r&proxyRequestFlagsIntermediate != 0 {
+		flags = append(flags, "INTERMEDIATE")
+	}
+	if r&proxyRequestFlagsAbdridged != 0 {
+		flags = append(flags, "ABRIDGED")
+	}
+	if r&proxyRequestFlagsQuickAck != 0 {
+		flags = append(flags, "QUICK_ACK")
+	}
+	if r&proxyRequestFlagsPad != 0 {
+		flags = append(flags, "PAD")
+	}
+
+	return strings.Join(flags, " | ")
+}

telegram/middle.go

@@ -0,0 +1,56 @@
+package telegram
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/telegram/api"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type middleTelegram struct {
+	baseTelegram
+
+	secret []byte
+	mutex  sync.RWMutex
+}
+
+func (m *middleTelegram) update() error {
+	secret, err := api.Secret()
+	if err != nil {
+		return fmt.Errorf("cannot fetch secret: %w", err)
+	}
+
+	v4Addresses, v4DefaultDC, err := api.AddressesV4()
+	if err != nil {
+		return fmt.Errorf("cannot fetch addresses for ipv4: %w", err)
+	}
+
+	v6Addresses, v6DefaultDC, err := api.AddressesV6()
+	if err != nil {
+		return fmt.Errorf("cannot fetch addresses for ipv6: %w", err)
+	}
+
+	m.mutex.Lock()
+	m.secret = secret
+	m.v4DefaultDC = v4DefaultDC
+	m.V6DefaultDC = v6DefaultDC
+	m.v4Addresses = v4Addresses
+	m.v6Addresses = v6Addresses
+	m.mutex.Unlock()
+
+	return nil
+}
+
+func (m *middleTelegram) Dial(ctx context.Context,
+	cancel context.CancelFunc,
+	dc conntypes.DC,
+	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
+	if dc == 0 {
+		dc = conntypes.DCDefaultIdx
+	}
+
+	return m.baseTelegram.dial(ctx, cancel, dc, protocol)
+}

wrappers/blockcipher.go

@@ -0,0 +1,157 @@
+package wrappers
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"errors"
+	"fmt"
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4
+
+type wrapperBlockCipher struct {
+	buf bytes.Buffer
+
+	parent    StreamReadWriteCloser
+	encryptor cipher.BlockMode
+	decryptor cipher.BlockMode
+}
+
+func (w *wrapperBlockCipher) Write(p []byte) (int, error) {
+	encrypted, err := w.encrypt(p)
+	if err != nil {
+		return 0, err
+	}
+	return w.parent.Write(encrypted)
+}
+
+func (w *wrapperBlockCipher) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	encrypted, err := w.encrypt(p)
+	if err != nil {
+		return 0, err
+	}
+	return w.parent.WriteTimeout(encrypted, timeout)
+}
+
+func (w *wrapperBlockCipher) Read(p []byte) (int, error) {
+	return w.read(p, readAll)
+
+}
+
+func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.read(p, readAllTimeout(timeout))
+}
+
+func (w *wrapperBlockCipher) read(p []byte, reader func(StreamReadWriteCloser) ([]byte, error)) (int, error) {
+	if w.buf.Len() > 0 {
+		return w.flush(p)
+	}
+
+	var buf []byte
+	for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
+		rv, err := reader(w.parent)
+		if err != nil {
+			return 0, fmt.Errorf("cannot read from socket: %w", err)
+		}
+		buf = append(buf, rv...)
+	}
+
+	w.decryptor.CryptBlocks(buf, buf)
+	w.buf.Write(buf)
+
+	return w.flush(p)
+}
+
+func (w *wrapperBlockCipher) flush(p []byte) (int, error) {
+	if w.buf.Len() > len(p) {
+		return w.buf.Read(p)
+	}
+
+	sizeToReturn := w.buf.Len()
+	copy(p, w.buf.Bytes())
+	w.buf.Reset()
+
+	return sizeToReturn, nil
+}
+
+func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) {
+	if len(p)%aes.BlockSize > 0 {
+		return nil, fmt.Errorf("incorrect block size %d", len(p))
+	}
+
+	encrypted := make([]byte, len(p))
+	w.encryptor.CryptBlocks(encrypted, p)
+
+	return encrypted, nil
+}
+
+func readAll(src StreamReadWriteCloser) (rv []byte, err error) {
+	buf := make([]byte, blockCipherReadCurrentDataBufferSize)
+	n := blockCipherReadCurrentDataBufferSize
+
+	for n == len(buf) {
+		n, err = src.Read(buf)
+		if err != nil {
+			return nil, err
+		}
+		rv = append(rv, buf[:n]...)
+	}
+
+	return rv, nil
+}
+
+func readAllTimeout(timeout time.Duration) func(StreamReadWriteCloser) ([]byte, error) {
+	return func(src StreamReadWriteCloser) (rv []byte, err error) {
+		tmo := timeout
+		buf := make([]byte, blockCipherReadCurrentDataBufferSize)
+		n := blockCipherReadCurrentDataBufferSize
+
+		for n == len(buf) {
+			if tmo <= 0 {
+				return nil, errors.New("timeout")
+			}
+			startTime := time.Now()
+			n, err = src.ReadTimeout(buf, tmo)
+			if err != nil {
+				return nil, err
+			}
+			rv = append(rv, buf[:n]...)
+			tmo -= time.Since(startTime)
+		}
+
+		return rv, nil
+	}
+}
+
+func (w *wrapperBlockCipher) Close() error {
+	return w.parent.Close()
+}
+
+func (w *wrapperBlockCipher) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperBlockCipher) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("block-cipher")
+}
+
+func (w *wrapperBlockCipher) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperBlockCipher) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func NewBlockCipher(parent StreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) StreamReadWriteCloser {
+	return &wrapperBlockCipher{
+		parent:    parent,
+		encryptor: encryptor,
+		decryptor: decryptor,
+	}
+}

wrappers/mtproto_frame.go

@@ -0,0 +1,159 @@
+package wrappers
+
+import (
+	"bytes"
+	"crypto/aes"
+	"encoding/binary"
+	"fmt"
+	"hash/crc32"
+	"io"
+	"io/ioutil"
+	"net"
+
+	"go.uber.org/zap"
+)
+
+const (
+	mtprotoFrameMinMessageLength = 12
+	mtprotoFrameMaxMessageLength = 16777216
+)
+
+var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00}
+
+// MTProtoFrame is a wrapper which converts written data to the MTProtoFrame.
+// The format of the frame:
+//
+// [ MSGLEN(4) | SEQNO(4) | MSG(...) | CRC32(4) | PADDING(4*x) ]
+//
+// MSGLEN is the length of the message + len of seqno and msglen.
+// SEQNO is the number of frame in the receive/send sequence. If client
+//   sends a message with SeqNo 18, it has to receive message with SeqNo 18.
+// MSG is the data which has to be written
+// CRC32 is the CRC32 checksum of MSGLEN + SEQNO + MSG
+// PADDING is custom padding schema to complete frame length to such that
+//    len(frame) % 16 == 0
+type wrapperMtprotoFrame struct {
+	parent     StreamReadWriteCloser
+	logger     *zap.SugaredLogger
+	readSeqNo  int32
+	writeSeqNo int32
+}
+
+func (w *wrapperMtprotoFrame) Read() (Packet, error) {
+	buf := &bytes.Buffer{}
+	sum := crc32.NewIEEE()
+	writer := io.MultiWriter(buf, sum)
+
+	for {
+		buf.Reset()
+		sum.Reset()
+		if _, err := io.CopyN(writer, w.parent, 4); err != nil {
+			return nil, fmt.Errorf("cannot read frame padding: %w", err)
+		}
+		if !bytes.Equal(buf.Bytes(), mtprotoFramePadding) {
+			break
+		}
+	}
+
+	messageLength := binary.LittleEndian.Uint32(buf.Bytes())
+	w.logger.Debugw("Read MTProto frame",
+		"messageLength", messageLength,
+		"sequence_number", w.readSeqNo,
+	)
+	if messageLength%4 != 0 || messageLength < mtprotoFrameMinMessageLength ||
+		messageLength > mtprotoFrameMaxMessageLength {
+		return nil, fmt.Errorf("Incorrect frame message length %d", messageLength)
+	}
+
+	buf.Reset()
+	buf.Grow(int(messageLength) - 4 - 4)
+	if _, err := io.CopyN(writer, w.parent, int64(messageLength)-4-4); err != nil {
+		return nil, fmt.Errorf("cannot read the message frame: %w", err)
+	}
+
+	var seqNo int32
+	binary.Read(buf, binary.LittleEndian, &seqNo) // nolint: errcheck, gosec
+	if seqNo != w.readSeqNo {
+		return nil, fmt.Errorf("unexpected sequence number %d (wait for %d)", seqNo, w.readSeqNo)
+	}
+
+	data, _ := ioutil.ReadAll(buf) // nolint: gosec
+	buf.Reset()
+	// write to buf, not to writer. This is because we are going to fetch
+	// crc32 checksum.
+	if _, err := io.CopyN(buf, w.parent, 4); err != nil {
+		return nil, fmt.Errorf("cannot read checksum: %w", err)
+	}
+
+	checksum := binary.LittleEndian.Uint32(buf.Bytes())
+	if checksum != sum.Sum32() {
+		return nil, fmt.Errorf("CRC32 checksum mismatch. wait for %d, got %d", sum.Sum32(), checksum)
+	}
+
+	w.logger.Debugw("Read MTProto frame",
+		"messageLength", messageLength,
+		"sequence_number", w.readSeqNo,
+		"dataLength", len(data),
+		"checksum", checksum,
+	)
+	w.readSeqNo++
+
+	return data, nil
+}
+
+func (w *wrapperMtprotoFrame) Write(p Packet) 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, w.writeSeqNo)
+	buf.Write(p)
+
+	checksum := crc32.ChecksumIEEE(buf.Bytes())
+	binary.Write(buf, binary.LittleEndian, checksum)
+	buf.Write(bytes.Repeat(mtprotoFramePadding, paddingLength/4))
+
+	w.logger.Debugw("Write MTProto frame",
+		"length", len(p),
+		"sequence_number", w.writeSeqNo,
+		"crc32", checksum,
+		"frame_length", buf.Len(),
+	)
+	w.writeSeqNo++
+
+	_, err := w.parent.Write(buf.Bytes())
+
+	return err
+}
+
+func (w *wrapperMtprotoFrame) Close() error {
+	return w.parent.Close()
+}
+
+func (w *wrapperMtprotoFrame) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperMtprotoFrame) Logger() *zap.SugaredLogger {
+	return w.logger
+}
+
+func (w *wrapperMtprotoFrame) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperMtprotoFrame) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func NewMtprotoFrame(parent StreamReadWriteCloser, seqNo int32) PacketReadWriteCloser {
+	return &wrapperMtprotoFrame{
+		parent:     parent,
+		logger:     parent.Logger().Named("mtproto-frame"),
+		readSeqNo:  seqNo,
+		writeSeqNo: seqNo,
+	}
+}