Упражнение: Разбор Protobuf

В этом упражнении вы напишете парсер для бинарного кодирования protobuf. Не волнуйтесь, это проще, чем кажется! Оно иллюстрирует распространённый паттерн разбора: передачу срезов данных. Сами базовые данные никогда не копируются.

Чтобы полностью разобрать сообщение protobuf, нужно знать типы полей, индексированные по их номерам. Обычно это задаётся в файле proto. В этом упражнении мы закодируем эту информацию в виде операторов match в функциях, которые вызываются для каждого поля.

Мы будем использовать следующий proto:

message PhoneNumber {
  optional string number = 1;
  optional string type = 2;
}

message Person {
  optional string name = 1;
  optional int32 id = 2;
  repeated PhoneNumber phones = 3;
}

Сообщения

Сообщение proto кодируется как последовательность полей, одно за другим. Каждое поле реализовано как «тег», за которым следует значение. Тег содержит номер поля (например, 2 для поля id сообщения Person) и тип передачи (wire type), определяющий, как по потоку байтов находить полезную нагрузку. Они объединяются в одно целое число, которое декодируется в unpack_tag ниже.

Varint

Целые числа, включая тег, представлены с помощью кодирования переменной длины, которое называется VARINT. К счастью, parse_varint уже определена ниже.

Типы передачи (Wire Types)

Proto определяет несколько типов передачи, из которых в этом упражнении используются только два.

Тип передачи Varint содержит одно varint и используется для кодирования значений proto типа int32, например Person.id.

Тип передачи Len содержит длину, выраженную как varint, за которой следует полезная нагрузка из стольких байтов. Он используется для кодирования значений proto типа string, например Person.name. Также он используется для значений proto, содержащих вложенные сообщения, например Person.phones, где полезная нагрузка содержит кодирование вложенного сообщения.

Упражнение

Приведённый код также определяет обратные вызовы для обработки полей Person и PhoneNumber, а также для разбора сообщения на последовательность вызовов этих обратных вызовов.

От вас остаётся реализовать функцию parse_field и трейт ProtoMessage для Person и PhoneNumber.

// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0

/// Тип передачи в том виде, в котором он появляется в потоке.
enum WireType {
    /// Тип передачи Varint указывает, что значение --- одно VARINT.
    Varint,
    // Тип передачи I64 указывает, что значение --- ровно 8 байт в порядке
    // little-endian, содержащие 64-битное знаковое целое или значение типа double.
    //I64,  -- не нужно для этого упражнения
    /// Тип передачи Len указывает, что значение --- длина, представленная как
    /// VARINT, за которой следует ровно столько байтов.
    Len,
    // Тип передачи I32 указывает, что значение --- ровно 4 байта в порядке
    // little-endian, содержащие 32-битное знаковое целое или значение типа float.
    //I32,  -- не нужно для этого упражнения
}

#[derive(Debug)]
/// Значение поля, типизированное на основе типа передачи.
enum FieldValue<'a> {
    Varint(u64),
    //I64(i64),  -- не нужно для этого упражнения
    Len(&'a [u8]),
    //I32(i32),  -- не нужно для этого упражнения
}

#[derive(Debug)]
/// Поле, содержащее номер поля и его значение.
struct Field<'a> {
    field_num: u64,
    value: FieldValue<'a>,
}

trait ProtoMessage<'a>: Default {
    fn add_field(&mut self, field: Field<'a>);
}

impl From<u64> for WireType {
    fn from(value: u64) -> Self {
        match value {
            0 => WireType::Varint,
            //1 => WireType::I64,  -- не нужно для этого упражнения
            2 => WireType::Len,
            //5 => WireType::I32,  -- не нужно для этого упражнения
            _ => panic!("Invalid wire type: {value}"),
        }
    }
}

impl<'a> FieldValue<'a> {
    fn as_str(&self) -> &'a str {
        let FieldValue::Len(data) = self else {
            panic!("Expected string to be a `Len` field");
        };
        std::str::from_utf8(data).expect("Invalid string")
    }

    fn as_bytes(&self) -> &'a [u8] {
        let FieldValue::Len(data) = self else {
            panic!("Expected bytes to be a `Len` field");
        };
        data
    }

    fn as_u64(&self) -> u64 {
        let FieldValue::Varint(value) = self else {
            panic!("Expected `u64` to be a `Varint` field");
        };
        *value
    }
}

/// Разбирает VARINT, возвращая разобранное значение и оставшиеся байты.
fn parse_varint(data: &[u8]) -> (u64, &[u8]) {
    for i in 0..7 {
        let Some(b) = data.get(i) else {
            panic!("Not enough bytes for varint");
        };
        if b & 0x80 == 0 {
            // Это последний байт VARINT, поэтому преобразуем его в
            // u64 и вернём.
            let mut value = 0u64;
            for b in data[..=i].iter().rev() {
                value = (value << 7) | (b & 0x7f) as u64;
            }
            return (value, &data[i + 1..]);
        }
    }

    // Более 7 байтов недопустимо.
    panic!("Too many bytes for varint");
}

/// Преобразует тег в номер поля и WireType.
fn unpack_tag(tag: u64) -> (u64, WireType) {
    let field_num = tag >> 3;
    let wire_type = WireType::from(tag & 0x7);
    (field_num, wire_type)
}


/// Разбирает поле, возвращая оставшиеся байты
fn parse_field(data: &[u8]) -> (Field<'_>, &[u8]) {
    let (tag, remainder) = parse_varint(data);
    let (field_num, wire_type) = unpack_tag(tag);
    let (fieldvalue, remainder) = match wire_type {
        _ => todo!("Based on the wire type, build a Field, consuming as many bytes as necessary.")
    };
    todo!("Return the field, and any un-consumed bytes.")
}

/// Разбирает сообщение в заданных данных, вызывая `T::add_field` для каждого поля в
/// сообщении.
///
/// Весь входной набор данных потребляется.
fn parse_message<'a, T: ProtoMessage<'a>>(mut data: &'a [u8]) -> T {
    let mut result = T::default();
    while !data.is_empty() {
        let parsed = parse_field(data);
        result.add_field(parsed.0);
        data = parsed.1;
    }
    result
}

#[derive(Debug, Default)]
struct PhoneNumber<'a> {
    number: &'a str,
    type_: &'a str,
}

#[derive(Debug, Default)]
struct Person<'a> {
    name: &'a str,
    id: u64,
    phone: Vec<PhoneNumber<'a>>,
}

// TODO: Реализуйте ProtoMessage для Person и PhoneNumber.

#[test]
fn test_id() {
    let person_id: Person = parse_message(&[0x10, 0x2a]);
    assert_eq!(person_id, Person { name: "", id: 42, phone: vec![] });
}

#[test]
fn test_name() {
    let person_name: Person = parse_message(&[
        0x0a, 0x0e, 0x62, 0x65, 0x61, 0x75, 0x74, 0x69, 0x66, 0x75, 0x6c, 0x20,
        0x6e, 0x61, 0x6d, 0x65,
    ]);
    assert_eq!(person_name, Person { name: "beautiful name", id: 0, phone: vec![] });
}

#[test]
fn test_just_person() {
    let person_name_id: Person =
        parse_message(&[0x0a, 0x04, 0x45, 0x76, 0x61, 0x6e, 0x10, 0x16]);
    assert_eq!(person_name_id, Person { name: "Evan", id: 22, phone: vec![] });
}

#[test]
fn test_phone() {
    let phone: Person = parse_message(&[
        0x0a, 0x00, 0x10, 0x00, 0x1a, 0x16, 0x0a, 0x0e, 0x2b, 0x31, 0x32, 0x33,
        0x34, 0x2d, 0x37, 0x37, 0x37, 0x2d, 0x39, 0x30, 0x39, 0x30, 0x12, 0x04,
        0x68, 0x6f, 0x6d, 0x65,
    ]);
    assert_eq!(
        phone,
        Person {
            name: "",
            id: 0,
            phone: vec![PhoneNumber { number: "+1234-777-9090", type_: "home" },],
        }
    );
}

// Соберите всё это в один разбор.
#[test]
fn test_full_person() {
    let person: Person = parse_message(&[
        0x0a, 0x07, 0x6d, 0x61, 0x78, 0x77, 0x65, 0x6c, 0x6c, 0x10, 0x2a, 0x1a,
        0x16, 0x0a, 0x0e, 0x2b, 0x31, 0x32, 0x30, 0x32, 0x2d, 0x35, 0x35, 0x35,
        0x2d, 0x31, 0x32, 0x31, 0x32, 0x12, 0x04, 0x68, 0x6f, 0x6d, 0x65, 0x1a,
        0x18, 0x0a, 0x0e, 0x2b, 0x31, 0x38, 0x30, 0x30, 0x2d, 0x38, 0x36, 0x37,
        0x2d, 0x35, 0x33, 0x30, 0x38, 0x12, 0x06, 0x6d, 0x6f, 0x62, 0x69, 0x6c,
        0x65,
    ]);
    assert_eq!(
        person,
        Person {
            name: "maxwell",
            id: 42,
            phone: vec![
                PhoneNumber { number: "+1202-555-1212", type_: "home" },
                PhoneNumber { number: "+1800-867-5308", type_: "mobile" },
            ]
        }
    );
}