bililive_core/packet/
mod.rs

1//! Packet types.
2
3use std::convert::TryInto;
4use std::io::{Cursor, Read, Write};
5
6use flate2::read::ZlibDecoder;
7use flate2::write::ZlibEncoder;
8use flate2::Compression;
9use nom::Err;
10use serde::Deserialize;
11use serde_json::json;
12
13pub use types::*;
14
15use crate::config::StreamConfig;
16use crate::errors::{IncompleteResult, ParseError};
17
18mod parser;
19mod types;
20
21#[cfg(test)]
22mod tests;
23
24type Result<T, E = ParseError> = std::result::Result<T, E>;
25
26/// Bililive packet.
27///
28/// Packet can be used to encode/parse raw bilibili live packets, and extract information from it.
29#[derive(Debug, Clone, Eq, PartialEq, Hash)]
30pub struct Packet {
31    packet_length: u32,
32    header_length: u16,
33    protocol_version: Protocol,
34    op: Operation,
35    seq_id: u32,
36    data: Vec<u8>,
37}
38
39impl Packet {
40    /// Set the protocol version.
41    pub fn set_proto(&mut self, protocol_version: Protocol) {
42        self.protocol_version = protocol_version;
43    }
44    /// Set the operation.
45    pub fn set_op(&mut self, op: Operation) {
46        self.op = op;
47    }
48    /// Set the sequence id. By default it's 1.
49    pub fn set_seq_id(&mut self, seq_id: u32) {
50        self.seq_id = seq_id;
51    }
52    /// Set the packet body.
53    /// Packet length will be updated automatically.
54    pub fn set_data<T: Into<Vec<u8>>>(&mut self, data: T) {
55        self.data = data.into();
56        self.packet_length = self.header_length as u32 + self.data.len() as u32;
57    }
58}
59
60impl Packet {
61    /// Construct a new packet.
62    ///
63    /// To construct a zlib-compressed packet, you should create a JSON/Int32BE packet first,
64    /// then call [`Packet::compress`](Packet::compress) to convert it to a zlib one.
65    pub fn new<T: Into<Vec<u8>>>(op: Operation, protocol_version: Protocol, data: T) -> Self {
66        let data = data.into();
67
68        Self {
69            packet_length: data.len() as u32 + 16,
70            header_length: 16,
71            protocol_version,
72            op,
73            seq_id: 1,
74            data,
75        }
76    }
77
78    /// Convert a JSON/Int32BE packet to a zlib-compressed one.
79    ///
80    /// # Errors
81    /// Return errors if compression fails.
82    pub fn compress(self) -> Result<Self> {
83        let raw = self.encode();
84
85        let mut z = ZlibEncoder::new(Vec::new(), Compression::default());
86        z.write_all(&raw)?;
87        let data = z.finish()?;
88
89        Ok(Self::new(self.op, Protocol::Zlib, data))
90    }
91}
92
93impl Packet {
94    #[allow(clippy::missing_panics_doc)]
95    #[must_use]
96    pub fn new_room_enter(config: &StreamConfig) -> Self {
97        Self::new(
98            Operation::RoomEnter,
99            Protocol::Json,
100            serde_json::to_vec(&json!({
101                "uid": config.uid(),
102                "roomid": config.room_id(),
103                "protover": 2,
104                "platform": "web",
105                "clientver": "1.8.2",
106                "type": 2,
107                "key": config.token()
108            }))
109            .unwrap(),
110        )
111    }
112}
113
114impl Packet {
115    /// Get the packet length.
116    #[must_use]
117    pub const fn packet_length(&self) -> u32 {
118        self.packet_length
119    }
120    /// Get the header length.
121    #[must_use]
122    pub const fn header_length(&self) -> u16 {
123        self.header_length
124    }
125    /// Get the sequence id.
126    #[must_use]
127    pub const fn seq_id(&self) -> u32 {
128        self.seq_id
129    }
130    /// Get the operation.
131    #[must_use]
132    pub const fn op(&self) -> Operation {
133        self.op
134    }
135    /// Get the protocol version.
136    #[must_use]
137    pub const fn proto(&self) -> Protocol {
138        self.protocol_version
139    }
140    /// Get bytes of the body.
141    #[must_use]
142    pub fn bytes(&self) -> &[u8] {
143        &self.data
144    }
145    /// Try to parse the body by json.
146    ///
147    /// # Errors
148    /// It may fail if the model is incorrect or it's not a json packet.
149    /// You may check the type of the packet by [`Packet::proto`](Packet::proto).
150    pub fn json<'a, T: Deserialize<'a>>(&'a self) -> Result<T> {
151        serde_json::from_slice(&self.data).map_err(ParseError::Json)
152    }
153    /// Try to parse the body by big endian int32.
154    ///
155    /// # Errors
156    /// It may fail if it's not a int packet.
157    /// You may check the type of the packet by [`Packet::proto`](Packet::proto).
158    pub fn int32_be(&self) -> Result<i32> {
159        Ok(i32::from_be_bytes(
160            self.data
161                .as_slice()
162                .try_into()
163                .map_err(|_| ParseError::Int32BE)?,
164        ))
165    }
166}
167
168impl Packet {
169    /// Encode the packet into bytes ready to be sent to the server.
170    #[must_use]
171    pub fn encode(&self) -> Vec<u8> {
172        let mut buf = Vec::with_capacity(self.packet_length as usize);
173        buf.extend(self.packet_length.to_be_bytes());
174        buf.extend(self.header_length.to_be_bytes());
175        buf.extend((self.protocol_version as u16).to_be_bytes());
176        buf.extend((self.op as u32).to_be_bytes());
177        buf.extend(self.seq_id.to_be_bytes());
178        buf.extend(&self.data);
179        buf
180    }
181
182    /// Parse the packet received from Bilibili server.
183    #[must_use]
184    pub fn parse(input: &[u8]) -> IncompleteResult<(&[u8], Self)> {
185        match parser::parse(input) {
186            Ok((input, packet)) => {
187                if packet.protocol_version == Protocol::Zlib {
188                    let mut z = ZlibDecoder::new(Cursor::new(packet.data));
189                    let mut buf = Vec::new();
190                    if let Err(e) = z.read_to_end(&mut buf) {
191                        return IncompleteResult::Err(ParseError::ZlibError(e));
192                    }
193
194                    match parser::parse(&buf) {
195                        Ok((_, packet)) => IncompleteResult::Ok((input, packet)),
196                        Err(Err::Incomplete(needed)) => {
197                            IncompleteResult::Err(ParseError::PacketError(format!(
198                                "incomplete buffer: {:?} needed",
199                                needed
200                            )))
201                        }
202                        Err(Err::Error(e) | Err::Failure(e)) => {
203                            IncompleteResult::Err(ParseError::PacketError(format!("{:?}", e)))
204                        }
205                    }
206                } else {
207                    IncompleteResult::Ok((input, packet))
208                }
209            }
210            Err(Err::Incomplete(needed)) => IncompleteResult::Incomplete(needed),
211            Err(Err::Error(e) | Err::Failure(e)) => {
212                IncompleteResult::Err(ParseError::PacketError(format!("{:?}", e)))
213            }
214        }
215    }
216}