Telegram на go, часть 2: бинарный протокол

В предыдущей части были описаны подходы, примененные при написании парсера для схемы MTProto. Статья получилась чуть более общей, чем я рассчитывал, на этот раз я постараюсь рассказать больше про специфику Telegram.

Клиент на Go продолжает развиваться, а мы вернемся в прошлое и вспомним, как писался для него сериализатор и десериализатор протокола.

Основы

Есть два пути десериализации: поточный и работающий с буфером. На практике в MTProto сообщение больше мегабайта передать не получится, поэтому я выбрал вариант с буфером: допустим, что мы всегда сможем держать полное сообщение в памяти.

Получится вот такая структура:

// Buffer implements low level binary (de-)serialization for TL.
type Buffer struct {
    Buf []byte
}

И еще, MTProto в основном выравнивает значения по 4 байтам (32 битам), вынесем это в константу:

// Word represents 4-byte sequence.
// Values in TL are generally aligned to Word.
const Word = 4

Сериализация

Зная, что почти всё в MTProto little-endian, мы можем начать с сериализации uint32:

// PutUint32 serializes unsigned 32-bit integer.
func (b *Buffer) PutUint32(v uint32) {
    t := make([]byte, Word)
    binary.LittleEndian.PutUint32(t, v)
    b.Buf = append(b.Buf, t...)
}

Все остальные значения мы будем сериализовать аналогично: сначала аллоцировали слайс (компилятор Go достаточно умный, чтобы в этом случае не помещать его в heap, так как размер слайса маленький и константа), потом записали туда значение, а затем дописали слайс в буфер.

Основная сложность заключается в том, чтобы аккуратно обработать все типы, так как подробных примеров в документации нет. Приходится искать примеры по тестам у других клиентов, например grammers, замечательного клиента на Rust для Telegram.

Десериализация

Возможно, не стоило использовать одну и ту же структуру для десериализации и сериализации, но на практике это оказалось удобным, а сам пакет gotd/td/bin получился довольно компактным.

Продолжим пример с uint32:

// Uint32 decodes unsigned 32-bit integer from Buffer.
func (b *Buffer) Uint32() (uint32, error) {
    if len(b.Buf) < Word {
        return 0, io.ErrUnexpectedEOF
    }
    v := binary.LittleEndian.Uint32(b.Buf)
    b.Buf = b.Buf[Word:]
    return v, nil
}

Сначала мы проверяем, что в буфере достаточно данных, а иначе возвращаем ошибку io.ErrUnexpectedEOF. Затем вычитываем значение и сдвигаем буфер на длину этого значения. Этот подход тоже будет общим для десериализации.

Строки

Строки ([]byte и string) сериализуются по-разному в зависимости от их длины и выравниваются по 4 байтам.

Если длина меньше или равна 253, то она пишется в первый байт, а затем строка:

b = append(b, byte(l))
b = append(b, v...)
currentLen := l + 1
// Padding:
b = append(b, make([]byte, nearestPaddedValueLength(currentLen)-currentLen)...)
return b

Если длина больше, то первым байтом будет 254, потом три байта длины в little-endian, а затем уже строка:

b = append(b, 254, byte(l), byte(l>>8), byte(l>>16))
b = append(b, v...)
currentLen := l + 4
// Padding:
b = append(b, make([]byte, nearestPaddedValueLength(currentLen)-currentLen)...)

Сигнатура функции encodeString(b []byte, v string) []byte позволяет просто дописывать в b, не аллоцируя каждый раз новый слайс:

// PutString serializes bare string.
func (b *Buffer) PutString(s string) {
    b.Buf = encodeString(b.Buf, s)
}

Кстати, в своем десериализаторе строк я нашел несколько багов с помощью фаззинга. Про него я писал в первой части, повторяться не буду.

Структуры

Мы базово разобрали, как происходит работа со значениями, но нам нужно сериализировать сложные типы. Делается это довольно просто: сначала пишется ID типа (это те самые #5b38c6c1 из схемы, просто как uint32), а затем все его поля по порядку, как в схеме.

Например, у нас есть такая структура (это упрощенный выдуманный пример из теста для сериализатора):

// msg#9bdd8f1a code:int32 message:string = Message;
type Message struct {
    Code    int32
    Message string
}

И сериализовать её c помощью нашего Buffer довольно просто:

// EncodeTo implements bin.Encoder.
func (m Message) Encode(b *Buffer) error {
    b.PutID(0x9bdd8f1a)
    b.PutInt32(m.Code)
    b.PutString(m.Message)
    return nil
}

Если попробовать вызвать Encode, в буфере будет следующее содержимое:

m := Message{
    Code:    204,
    Message: "Wake up, Neo",
}
b := new(Buffer)
_ = m.Encode(b)

raw := []byte{
    // Type ID.
    0x1a, 0x8f, 0xdd, 0x9b,

    // Code as int32.
    204, 0x00, 0x00, 0x00,

    // String length.
    byte(len(m.Message)),

    // "Wake up, Neo" in hex.
    0x57, 0x61, 0x6b,
    0x65, 0x20, 0x75, 0x70,
    0x2c, 0x20, 0x4e, 0x65,
    0x6f, 0x00, 0x00, 0x00,
}

Для десериализации мы должны сначала проверить, что у нас действительно нужный тип. Лучше это делать без изменения Buf, а значит нам потребуется новый метод:

// PeekID returns next type id in Buffer, but does not consume it.
func (b *Buffer) PeekID() (uint32, error) {
    if len(b.Buf) < Word {
        return 0, io.ErrUnexpectedEOF
    }
    v := binary.LittleEndian.Uint32(b.Buf)
    return v, nil
}

Я опущу тело метода ConsumeID(id uint32) для краткости: он вызывает PeekID и возвращет ошибку на неожиданный тип, сдвигая буфер только в случае успеха. С помощью него получится вот такая десериализация:

func (m *Message) Decode(b *Buffer) error {
    if err := b.ConsumeID(0x9bdd8f1a); err != nil {
        return err
    }
    {
        v, err := b.Int32()
        if err != nil {
            return err
        }
        m.Code = v
    }
    {
        v, err := b.String()
        if err != nil {
            return err
        }
        m.Message = v
    }

    return nil
}

Теперь мы можем аналогично написать (де-)сериализатор для любого типа из схемы, имплементируя следующие интерфейсы:

// Encoder can encode it's binary form to Buffer.
type Encoder interface {
    Encode(b *Buffer) error
}

// Decoder can decode it's binary form from Buffer.
type Decoder interface {
    Decode(b *Buffer) error
}

Специальные случаи

Некоторые нюансы бинарного протокола заслуживают отдельного рассмотрения, так как могут быть не очевидны из документации.

Массивы

Называются векторами. Например:

messageActionChatCreate#a6638b9a title:string users:Vector<int> = MessageAction;

Тут понятно, как сериализовать title, но что делать с users?

Непосредственно Vector определен вот так:

vector#0x1cb5c415 {t:Type} # [ t ] = Vector t

Но что это значит в плане бинарного представления с первого взгляда не совсем понятно. Да и со второго тоже, так как понять из схемы это невозможно, а описано как специальный случай.

На практике вектор сериализуется следующим образом: сначала пишется тип вектора (0x1cb5c415), затем количество элементов, а потом уже сами элементы подряд:

// PutVectorHeader serializes vector header with provided length.
func (b *Buffer) PutVectorHeader(length int) {
    b.PutID(TypeVector)
    b.PutInt32(int32(length))
}

Значит, для массива из 10 чисел uint32, мы вызовем PutVectorHeader(10), а затем запишем 10 раз uint32.

Булевые значения

С ними примерно так же, как и с векторами, но чуть проще:

boolTrue#997275b5 = Bool;
boolFalse#bc799737 = Bool;

Таким образом, тобы записать Bool, нужно написать либо 0x997275b5, либо 0xbc799737:

const (
    TypeTrue  = 0x997275b5 // boolTrue#997275b5 = Bool
    TypeFalse = 0xbc799737 // boolFalse#bc799737 = Bool    
)

// PutBool serializes bare boolean.
func (b *Buffer) PutBool(v bool) {
    switch v {
    case true:
        b.PutID(TypeTrue)
    case false:
        b.PutID(TypeFalse)
    }
}

Для булевых значений есть специальный случай, они могут передаваться битовой маской, если являются частью флагов, о них чуть позже.

Флаги и опциональные поля

Схема поддерживает возможность не передавать значение некоторых полей, если они не нужны, то есть сделать поля опциональными. Это похоже на битовое поле: если соответствующий бит задан, то поле (де-)сериализируется, если нет, то пропускается.

Особым случаем опционального поля являются булевые значения (flags.0?true): они целиком кодируются одним битом в маске, непосредственная сериализация 0x997275b5 не используется, так как избыточна.

Но есть нюанс! Если поле определено как flags.0?Bool, то оно сериализуется как обычный Bool, значит есть вариант опционального булевого поля, которое избыточно сериализуется. Видимо, это legacy.

Представить такой bitfield в Go можно довольно просто:

// Fields represent a bitfield value that compactly encodes
// information about provided conditional fields.
type Fields uint32

// Has reports whether field with index n was set.
func (f Fields) Has(n int) bool {
    return f&(1<<n) != 0
}

// Set sets field with index n.
func (f *Fields) Set(n int) {
    *f |= 1 << n
}

Это поле сериализуется как обычный uint32.

А непосредственное использование будет таким:

// msg flags:# escape:flags.0?true ttl_seconds:flags.1?int = Message;
type FieldsMessage struct {
    Flags      bin.Fields
    Escape     bool
    TTLSeconds int
}

func (f *FieldsMessage) Encode(b *bin.Buffer) error {
    b.PutID(FieldsMessageTypeID)
    if f.Escape {
        f.Flags.Set(0)
    }
    if err := f.Flags.Encode(b); err != nil {
        return err
    }
    if f.Flags.Has(1) {
        b.PutInt(f.TTLSeconds)
    }
    return nil
}

Видно, что TTLSeconds пишется только если задан флаг 1, а Escape полностью кодируется в Flags.

Криптографические ключи

Еще одим специальным случаем будут типы int128 и int256:

int128 4*[ int ] = Int128;
int256 8*[ int ] = Int256;

В go их можно представить как массивы:

type Int128 [16]byte
type Int256 [32]byte

Таким образом, сериализуются они очень просто:

func (b *Buffer) PutInt128(v Int128) {
    b.Buf = append(b.Buf, v[:]...)
}
func (b *Buffer) PutInt256(v Int256) {
    b.Buf = append(b.Buf, v[:]...)
}

Самым сложным было понять, как конвертировать значения между ними и big.Int.

MTProto использует big-endian для этих значений по причине аналогичного поведения в OpenSSL. В Go big.Int поступает так же.

Получается, можно просто использовать SetBytes для чтения и FillBytes для записи:

var v Int256

i := new(big.Int).SetBytes(v[:]) // v -> i
i.FillBytes(v[:]) // i -> v

Итог

Пакет bin теперь готов, и появилась возможность сериализовать любой тип. Проблема в том, что в схеме очень, очень много типов и писать их вручную будет крайне накладно.

Эта проблема решается генерацией кода (де-)сериализации из схемы (вот для этого мы и писали парсер!). Возможно, я уделю генератору отдельную часть в серии статей. Этот модуль проекта получился сложным, переписывался несколько раз и я бы хотел хоть немного облегчить жизнь людям, которые будут писать кодогенераторы на Go для других форматов.

Для справки, на текущий момент генерируется около 180K SLOC из схем телеграма (api, mtproto, секретные чаты).

Хочу поблагодарить tdakkota и zweihander за их неоценимый вклад в развитие проекта! Без вас было бы очень тяжело.