source_demo_tool/
protobuf_message.rs

1use std::{io::Read, collections::BTreeMap};
2use buf_redux::BufReader;
3use source_demo_tool_impl_proc_macros::event;
4
5use crate::{parse_tools::{parse_varint, parse_u32, ParseVarIntExit}, event_data::EventData, demo_file::packet::protobuf_value::ProtobufValue};
6
7pub trait ProtobufMessageEnumTraits: Send + Sync {
8    fn to_vec(&self) -> Vec<(&'static str, ProtobufValue)>;
9    fn type_count() -> usize where Self: Sized;
10    fn to_str(&self) -> &'static str;
11    fn get_id_map() -> BTreeMap<usize, &'static str> where Self: Sized;
12    fn as_u64(&self) -> u64;
13}
14
15#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
16pub struct ProtobufMessage {
17    pub message: WireMessage,
18    pub field_number: u8,
19}
20
21#[derive(Debug)]
22pub enum FromReadableErr {
23    BufferErr(EventData),
24    InvalidOrCorrupt(EventData),
25    UnsupportedWireType(EventData, u8)
26}
27
28impl ProtobufMessage {
29    pub fn many_from_vec(raw_data: &Vec<u8>) -> Result<Vec<Self>, FromReadableErr> {
30        let mut reader = BufReader::with_capacity
31            (raw_data.len(), raw_data.as_slice());
32        reader.read_into_buf().unwrap();
33
34        let mut msgs = Vec::new();
35        while reader.buf_len() > 0 {
36            msgs.push(
37                match Self::from_readable(&mut reader) {
38                    Ok(val) => val,
39                    Err(e) => return Err(e)
40                }
41            );
42        }
43        Ok(msgs)
44    }
45
46    pub fn from_readable(mut reader: impl Read) -> Result<ProtobufMessage, FromReadableErr> {
47        let tag = match parse_varint(&mut reader) {
48            ParseVarIntExit::Ok(n) => n,
49            _ => return Err(FromReadableErr::InvalidOrCorrupt(event!{"bad tag varint read"}))
50        };
51
52        let message = match (tag & 0b111) as u8 {
53            wire_type::VARINT => WireMessage::VarInt(match parse_varint(&mut reader) {
54                ParseVarIntExit::Ok(n) => n,
55                _ => return Err(FromReadableErr::InvalidOrCorrupt(event!{"bad varint read, message"}))
56            }),
57            wire_type::LENGTH => WireMessage::Length({
58                let len = match parse_varint(&mut reader) {
59                    ParseVarIntExit::Ok(n) => n,
60                _ => return Err(FromReadableErr::InvalidOrCorrupt(event!{"bad varint read, len"}))
61                };
62
63                let mut v = Vec::new();
64                v.resize(len as usize, 0);
65                if let Err(_) = reader.read_exact(v.as_mut_slice()) {
66                    return Err(FromReadableErr::BufferErr(event!{"length data read"}))
67                }
68                v
69            }),
70            wire_type::FIXED32 => WireMessage::Fixed32(match parse_u32(&mut reader) {
71                Ok(n) => n,
72                Err(e) => {
73                    let mut ev = event!{""};
74                    ev.details = e.into();
75                    return Err(FromReadableErr::InvalidOrCorrupt(ev))
76                }
77            }),
78            n => return Err(FromReadableErr::UnsupportedWireType(event!{"bad wire type"}, n))
79        };
80
81        let field_number = ((tag & !0b111) >> 3) as u8;
82
83        Ok(ProtobufMessage { message, field_number })
84    }
85}
86
87#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
88pub enum WireMessage {
89    VarInt(u64),
90    Length(Vec<u8>),
91    Fixed32(u32)
92}
93pub mod wire_type {
94    pub const VARINT: u8 = 0;
95    pub const LENGTH: u8 = 2;
96    pub const FIXED32: u8 = 5;
97}
source_demo_tool/protobuf_message.rs

source_demo_tool/
protobuf_message.rs
