tycho_util/
tl.rs

1use bytes::Bytes;
2use tl_proto::{TlError, TlPacket, TlRead, TlResult, TlWrite};
3
4pub mod signature_ref {
5    use super::*;
6
7    #[inline]
8    pub fn size_hint(signature: &[u8; 64]) -> usize {
9        signature.as_slice().max_size_hint()
10    }
11
12    #[inline]
13    pub fn write<P: TlPacket>(signature: &[u8; 64], packet: &mut P) {
14        signature.as_slice().write_to(packet);
15    }
16
17    pub fn read<'a>(packet: &mut &'a [u8]) -> TlResult<&'a [u8; 64]> {
18        <&tl_proto::BoundedBytes<64>>::read_from(packet)
19            .and_then(|bytes| bytes.as_ref().try_into().map_err(|_e| TlError::InvalidData))
20    }
21}
22
23pub mod signature_owned {
24    use super::*;
25
26    #[inline]
27    pub fn size_hint(signature: &[u8; 64]) -> usize {
28        signature.as_slice().max_size_hint()
29    }
30
31    #[inline]
32    pub fn write<P: TlPacket>(signature: &[u8; 64], packet: &mut P) {
33        signature.as_slice().write_to(packet);
34    }
35
36    pub fn read(packet: &mut &[u8]) -> TlResult<Box<[u8; 64]>> {
37        <&tl_proto::BoundedBytes<64>>::read_from(packet).and_then(|bytes| {
38            let Ok::<[u8; 64], _>(bytes) = bytes.as_ref().try_into() else {
39                return Err(TlError::InvalidData);
40            };
41            Ok(Box::new(bytes))
42        })
43    }
44}
45
46pub mod signature_arc {
47    use std::sync::Arc;
48
49    use super::*;
50
51    #[inline]
52    pub fn size_hint(signature: &[u8; 64]) -> usize {
53        signature.as_slice().max_size_hint()
54    }
55
56    #[inline]
57    pub fn write<P: TlPacket>(signature: &[u8; 64], packet: &mut P) {
58        signature.as_slice().write_to(packet);
59    }
60
61    pub fn read(packet: &mut &[u8]) -> TlResult<Arc<[u8; 64]>> {
62        <&tl_proto::BoundedBytes<64>>::read_from(packet).and_then(|bytes| {
63            let Ok::<[u8; 64], _>(bytes) = bytes.as_ref().try_into() else {
64                return Err(TlError::InvalidData);
65            };
66            Ok(Arc::new(bytes))
67        })
68    }
69}
70
71pub struct VecWithMaxLen<const N: usize>;
72
73impl<const N: usize> VecWithMaxLen<N> {
74    #[inline]
75    pub fn size_hint<T: tl_proto::TlWrite>(value: &[T]) -> usize {
76        value.max_size_hint()
77    }
78
79    #[inline]
80    pub fn write<P: TlPacket, T: tl_proto::TlWrite>(value: &[T], packet: &mut P) {
81        value.write_to(packet);
82    }
83
84    pub fn read<'tl, T>(packet: &mut &'tl [u8]) -> TlResult<Vec<T>>
85    where
86        T: tl_proto::TlRead<'tl>,
87    {
88        let len = u32::read_from(packet)? as usize;
89        if len > N {
90            return Err(TlError::InvalidData);
91        }
92
93        let mut items = Vec::with_capacity(len);
94        for _ in 0..len {
95            items.push(T::read_from(packet)?);
96        }
97
98        Ok(items)
99    }
100}
101
102pub struct BigBytes<const MAX_SIZE: usize>;
103
104impl<const MAX_SIZE: usize> BigBytes<MAX_SIZE> {
105    pub const MAX_SIZE: usize = MAX_SIZE;
106
107    #[inline]
108    pub fn size_hint<T: AsRef<[u8]>>(bytes: &T) -> usize {
109        BigBytesRef::<MAX_SIZE>::size_hint(bytes)
110    }
111
112    #[inline]
113    pub fn write<P: TlPacket>(bytes: &[u8], packet: &mut P) {
114        BigBytesRef::<MAX_SIZE>::write(bytes, packet);
115    }
116
117    #[inline]
118    pub fn read(packet: &mut &[u8]) -> TlResult<Bytes> {
119        BigBytesRef::<MAX_SIZE>::read(packet).map(Bytes::copy_from_slice)
120    }
121}
122
123pub struct BigBytesRef<const MAX_SIZE: usize>;
124
125impl<const MAX_SIZE: usize> BigBytesRef<MAX_SIZE> {
126    pub const MAX_SIZE: usize = MAX_SIZE;
127
128    pub fn size_hint<T: AsRef<[u8]>>(bytes: &T) -> usize {
129        let len = bytes.as_ref().len();
130        4 + len + big_bytes_padding(len)
131    }
132
133    pub fn write<P: TlPacket>(bytes: &[u8], packet: &mut P) {
134        const PADDING: [u8; 3] = [0; 3];
135
136        let len = bytes.len();
137        packet.write_u32(len as u32);
138        packet.write_raw_slice(bytes);
139        if !len.is_multiple_of(4) {
140            packet.write_raw_slice(&PADDING[0..4 - len % 4]);
141        }
142    }
143
144    pub fn read<'tl>(packet: &mut &'tl [u8]) -> TlResult<&'tl [u8]> {
145        let len = u32::read_from(packet)? as usize;
146        if len > Self::MAX_SIZE {
147            return Err(tl_proto::TlError::InvalidData);
148        }
149        let padding = big_bytes_padding(len);
150
151        if packet.len() < len + padding {
152            return Err(tl_proto::TlError::UnexpectedEof);
153        }
154
155        let bytes = &packet[..len];
156        *packet = &packet[len + padding..];
157
158        Ok(bytes)
159    }
160}
161
162const fn big_bytes_padding(len: usize) -> usize {
163    (4 - len % 4) % 4
164}
165
166#[cfg(test)]
167mod tests {
168    use bytes::Bytes;
169
170    use super::*;
171
172    #[test]
173    fn big_bytes() {
174        type BigEnough = BigBytes<{ 100 << 20 }>;
175
176        // For each padding
177        for i in 0..4 {
178            let big_bytes = Bytes::from(vec![123; 1000 + i]);
179
180            let mut serialized = Vec::new();
181            BigEnough::write(&big_bytes, &mut serialized);
182
183            // Must be aligned by 4
184            assert_eq!(serialized.len() % 4, 0);
185
186            let packet = &mut serialized.as_slice();
187            let deserialized = BigEnough::read(packet).unwrap();
188            // Must be equal
189            assert_eq!(big_bytes, deserialized);
190
191            // Must consume all bytes
192            assert!(packet.is_empty());
193        }
194    }
195}
tycho_util/tl.rs

tycho_util/
tl.rs
