1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 // http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13
14 use std::convert::TryFrom;
15 use std::io::Read;
16
17 pub type Result<T> = std::result::Result<T, Err>;
18
19 // pub trait TransferData {
20 // fn peek_u(&self, ty: &str, len: usize) -> &str;
21 // fn get_u(&self, ty: &str, len: usize) -> &str;
22 // fn put_u(&self, ty: &str, value: usize,len: usize) -> &mut [u8];
23 // }
24 //
25 // impl TransferData for Octets {
26 // fn peek_u(&self, ty: &str, len: usize) -> &str {
27 // todo!()
28 // }
29 //
30 // fn get_u(&self, ty: &str, len: usize) -> &str {
31 // todo!()
32 // }
33 //
34 // fn put_u(&self, ty: &str, value: usize, len: usize) -> &mut [u8] {
35 // todo!()
36 // }
37 // }
38 // buf fragment splited by out offset
39 #[derive(Debug, PartialEq, Eq)]
40 pub struct ReadVarint<'a> {
41 buf: &'a [u8],
42 }
43
44 impl<'a> ReadVarint<'a> {
new(buf: &'a [u8]) -> Self45 pub fn new(buf: &'a [u8]) -> Self {
46 ReadVarint { buf }
47 }
48
into_u8(&mut self) -> Result<u8>49 pub fn into_u8(&mut self) -> Result<u8> {
50 const len: usize = 1;
51
52 if self.buf.len() < len {
53 return Err(BufferTooShortError);
54 }
55 let bytes: [u8; len] = <[u8; len]>::try_from(self.buf[..len].as_ref()).unwrap();
56 if cfg!(target_endian = "big") {
57 let res: u8 = u8::from_be_bytes(bytes);
58 Ok(res)
59 } else {
60 let res: u8 = u8::from_le_bytes(bytes);
61 Ok(res)
62 }
63 }
64
into_u16(&mut self) -> Result<u16>65 pub fn into_u16(&mut self) -> Result<u16> {
66 const len: usize = 2;
67
68 if self.buf.len() < len {
69 return Err(BufferTooShortError);
70 }
71 let bytes: [u8; len] = <[u8; len]>::try_from(self.buf[..len].as_ref()).unwrap();
72 if cfg!(target_endian = "big") {
73 let res: u16 = u16::from_be_bytes(bytes);
74 Ok(res)
75 } else {
76 let res: u16 = u16::from_le_bytes(bytes);
77 Ok(res)
78 }
79 }
80
into_u32(&mut self) -> Result<u32>81 pub fn into_u32(&mut self) -> Result<u32> {
82 const len: usize = 4;
83
84 if self.buf.len() < len {
85 return Err(BufferTooShortError);
86 }
87 let bytes: [u8; len] = <[u8; len]>::try_from(self.buf[..len].as_ref()).unwrap();
88 if cfg!(target_endian = "big") {
89 let res: u32 = u32::from_be_bytes(bytes);
90 Ok(res)
91 } else {
92 let res: u32 = u32::from_le_bytes(bytes);
93 Ok(res)
94 }
95 }
96
into_u64(&mut self) -> Result<u64>97 pub fn into_u64(&mut self) -> Result<u64> {
98 const len: usize = 8;
99
100 if self.buf.len() < len {
101 return Err(BufferTooShortError);
102 }
103 let bytes: [u8; len] = <[u8; len]>::try_from(self.buf[..len].as_ref()).unwrap();
104 if cfg!(target_endian = "big") {
105 let res: u64 = u64::from_be_bytes(bytes);
106 Ok(res)
107 } else {
108 let res: u64 = u64::from_le_bytes(bytes);
109 Ok(res)
110 }
111 }
112
113 /// Reads an unsigned variable-length integer in network byte-order from
114 /// the current offset and advances the buffer.
get_varint(&mut self) -> Result<u64>115 pub fn get_varint(&mut self) -> Result<u64> {
116 let first = self.into_u8()?;
117
118 let len = varint_parse_len(first);
119
120 if len > self.cap() {
121 return Err(BufferTooShortError);
122 }
123
124 let out = match len {
125 1 => u64::from(self.into_u8()?),
126
127 2 => u64::from(self.into_u16()? & 0x3fff),
128
129 4 => u64::from(self.into_u32()? & 0x3fffffff),
130
131 8 => self.into_u64()? & 0x3fffffffffffffff,
132
133 _ => unreachable!(),
134 };
135
136 Ok(out)
137 }
138
139 /// Returns the remaining capacity in the buffer.
cap(&self) -> usize140 pub fn cap(&self) -> usize {
141 self.buf.len() - self.off
142 }
143 }
144
145 // Component encoding status.
146 enum TokenStatus<T, E> {
147 // The current component is completely encoded.
148 Complete(T),
149 // The current component is partially encoded.
150 Partial(E),
151 }
152
153 type TokenResult<T> = Result<TokenStatus<usize, T>>;
154
155 struct WriteData<'a> {
156 src: &'a [u8],
157 src_idx: &'a mut usize,
158 dst: &'a mut [u8],
159 }
160
161 impl<'a> WriteData<'a> {
new(src: &'a [u8], src_idx: &'a mut usize, dst: &'a mut [u8]) -> Self162 fn new(src: &'a [u8], src_idx: &'a mut usize, dst: &'a mut [u8]) -> Self {
163 WriteData { src, src_idx, dst }
164 }
165
write(&mut self) -> TokenResult<usize>166 fn write(&mut self) -> TokenResult<usize> {
167 let src_idx = *self.src_idx;
168 let input_len = self.src.len() - src_idx;
169 let output_len = self.dst.len();
170 let num = (&self.src[src_idx..]).read(self.dst).unwrap();
171 if output_len >= input_len {
172 return Ok(TokenStatus::Complete(num));
173 }
174 *self.src_idx += num;
175 Ok(TokenStatus::Partial(num))
176 }
177 }
178
179 pub struct WriteVarint<'a> {
180 src: &'a [u8],
181 src_idx: &'a mut usize,
182 dst: &'a mut [u8],
183 }
184
185 impl<'a> WriteVarint<'a> {
186 // pub fn new(buf: &'a mut [u8]) -> Self {
187 // WriteVarint { buf }
188 // }
new(src: &'a [u8], src_idx: &'a mut usize, dst: &'a mut [u8]) -> Self189 pub fn new(src: &'a [u8], src_idx: &'a mut usize, dst: &'a mut [u8]) -> Self {
190 // src需要从value转码过来
191 WriteVarint { src, src_idx, dst }
192 }
193
194 /// Writes an unsigned 8-bit integer at the current offset and advances
195 /// the buffer.
write_u8(&mut self, value: u8) -> Result<usize>196 pub fn write_u8(&mut self, value: u8) -> Result<usize> {
197 const len: usize = 1;
198 // buf长度不够问题返回err，再由外层处理
199 if self.buf.len() != len {
200 return Err(BufferTooShortError);
201 }
202
203 let bytes: [u8; len] = value.to_be_bytes();
204 self.buf.copy_from_slice(bytes.as_slice());
205 Ok(len)
206 }
207
write_u16(&mut self, value: u16) -> Result<usize>208 pub fn write_u16(&mut self, value: u16) -> Result<usize> {
209 const len: usize = 2;
210 // buf长度不够问题返回err，再由外层处理
211 if self.buf.len() != len {
212 return Err(BufferTooShortError);
213 }
214
215 let bytes: [u8; len] = value.to_be_bytes();
216 self.buf.copy_from_slice(bytes.as_slice());
217 Ok(len)
218 }
219
write_u32(&mut self, value: u32) -> Result<usize>220 pub fn write_u32(&mut self, value: u32) -> Result<usize> {
221 const len: usize = 4;
222 // buf长度不够问题返回err，再由外层处理
223 if self.buf.len() != len {
224 return Err(BufferTooShortError);
225 }
226
227 let bytes: [u8; len] = value.to_be_bytes();
228 self.buf.copy_from_slice(bytes.as_slice());
229 Ok(len)
230 }
231
write_u64(&mut self, value: u64) -> Result<usize>232 pub fn write_u64(&mut self, value: u64) -> Result<usize> {
233 const len: usize = 8;
234 // buf长度不够问题返回err，再由外层处理
235 if self.buf.len() != len {
236 return Err(BufferTooShortError);
237 }
238
239 let bytes: [u8; len] = value.to_be_bytes();
240 self.buf.copy_from_slice(bytes.as_slice());
241 Ok(len)
242 }
243
244 /// Writes an unsigned variable-length integer in network byte-order at the
245 /// current offset and advances the buffer.
write_varint(&mut self, value: u64) -> Result<usize>246 pub fn write_varint(&mut self, value: u64) -> Result<usize> {
247 self.write_varint_with_len(value, varint_len(value))
248 }
249
write_varint_with_len(&mut self, value: u64, len: usize) -> Result<usize>250 pub fn write_varint_with_len(&mut self, value: u64, len: usize) -> Result<usize> {
251 if self.cap() < len {
252 return Err(BufferTooShortError);
253 }
254
255 let res = match len {
256 1 => self.write_u8(value as u8)?,
257
258 2 => {
259 let size = self.write_u16(value as u16)?;
260 *self.buf[0] |= 0x40;
261 size
262 }
263
264 4 => {
265 let size = self.write_u32(value as u32)?;
266 *self.buf[0] |= 0x80;
267 size
268 }
269
270 8 => {
271 let size = self.write_u64(value)?;
272 *self.buf[0] |= 0xc0;
273 size
274 }
275
276 _ => panic!("value is too large for varint"),
277 };
278
279 Ok(res)
280 }
281 }
282
283 /// Returns how many bytes it would take to encode `v` as a variable-length
284 /// integer.
varint_len(v: u64) -> usize285 pub const fn varint_len(v: u64) -> usize {
286 match v {
287 0..=63 => 1,
288 64..=16383 => 2,
289 16384..=1_073_741_823 => 4,
290 1_073_741_824..=4_611_686_018_427_387_903 => 8,
291 _ => {unreachable!()}
292 }
293 }
294
295 /// Returns how long the variable-length integer is, given its first byte.
varint_parse_len(byte: u8) -> usize296 pub const fn varint_parse_len(byte: u8) -> usize {
297 let byte = byte >> 6;
298 if byte <= 3 {
299 1 << byte
300 } else {
301 unreachable!()
302 }
303 }
304