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 use ylong_runtime::iter::parallel::ParSplit;
18
19 use crate::h3::error::CommonError::BufferTooShort;
20 use crate::h3::error::{EncodeError, H3Error};
21
22 pub type Result<T> = std::result::Result<T, H3Error>;
23
24 macro_rules! peek_bytes {
25 ($buf: expr, $ty: ty, $len: expr) => {{
26 if $buf.len() < $len {
27 return Err(H3Error::Serialize(BufferTooShort));
28 }
29 let bytes: [u8; $len] = <[u8; $len]>::try_from($buf[..$len].as_ref()).unwrap();
30 let res = <$ty>::from_be_bytes(bytes);
31 Ok(res)
32 }};
33 }
34
35 macro_rules! poll_bytes {
36 ($this: expr, $ty: ty, $len: expr) => {{
37 let res = peek_bytes!($this.buf[$this.idx..], $ty, $len);
38 $this.idx += $len;
39 res
40 }};
41 }
42
43 #[derive(Debug, PartialEq, Eq)]
44 pub struct ReadableBytes<'a> {
45 buf: &'a [u8],
46 idx: usize,
47 }
48
49 impl<'a> ReadableBytes<'a> {
from(buf: &'a [u8]) -> Self50 pub fn from(buf: &'a [u8]) -> Self {
51 ReadableBytes { buf, idx: 0 }
52 }
53
peek_u8(&mut self) -> Result<u8>54 pub(crate) fn peek_u8(&mut self) -> Result<u8> {
55 peek_bytes!(self.buf[self.idx..], u8, 1)
56 }
57
poll_u8(&mut self) -> Result<u8>58 pub fn poll_u8(&mut self) -> Result<u8> {
59 poll_bytes!(self, u8, 1)
60 }
61
poll_u16(&mut self) -> Result<u16>62 pub fn poll_u16(&mut self) -> Result<u16> {
63 poll_bytes!(self, u16, 2)
64 }
65
poll_u32(&mut self) -> Result<u32>66 pub fn poll_u32(&mut self) -> Result<u32> {
67 poll_bytes!(self, u32, 4)
68 }
69
poll_u64(&mut self) -> Result<u64>70 pub fn poll_u64(&mut self) -> Result<u64> {
71 poll_bytes!(self, u64, 8)
72 }
73
74 /// Reads an unsigned variable-length integer in network byte-order from
75 /// the current offset and advances the buffer.
get_varint(&mut self) -> Result<u64>76 pub fn get_varint(&mut self) -> Result<u64> {
77 let first = self.peek_u8()?;
78 let len = parse_varint_len(first);
79 if len > self.cap() {
80 return Err(BufferTooShort.into());
81 }
82 let out = match len {
83 1 => u64::from(self.poll_u8()?),
84
85 2 => u64::from(self.poll_u16()? & 0x3fff),
86
87 4 => u64::from(self.poll_u32()? & 0x3fffffff),
88
89 8 => self.poll_u64()? & 0x3fffffffffffffff,
90
91 _ => unreachable!(),
92 };
93
94 Ok(out)
95 }
96
97 /// Returns the remaining capacity in the buffer.
cap(&self) -> usize98 pub fn cap(&self) -> usize {
99 self.buf.len() - self.idx
100 }
101
index(&self) -> usize102 pub fn index(&self) -> usize {
103 self.idx
104 }
105
remaining(&self) -> &[u8]106 pub fn remaining(&self) -> &[u8] {
107 &self.buf[self.idx..]
108 }
109
slice(&mut self, length: usize) -> Result<&[u8]>110 pub fn slice(&mut self, length: usize) -> Result<&[u8]> {
111 if self.cap() < length {
112 Err(BufferTooShort.into())
113 } else {
114 let curr = self.idx;
115 self.idx += length;
116 Ok(&self.buf[curr..self.idx])
117 }
118 }
119 }
120
121 macro_rules! write_bytes {
122 ($this: expr, $value: expr, $len: expr) => {{
123 // buf长度不够问题返回err，再由外层处理
124 if $this.remaining() < $len {
125 return Err(BufferTooShort.into());
126 }
127
128 let mut bytes: [u8; $len] = $value.to_be_bytes();
129 match $len {
130 1 => {}
131 2 => bytes[0] |= 0x40,
132 4 => bytes[0] |= 0x80,
133 8 => bytes[0] |= 0xC0,
134 _ => unreachable!(),
135 }
136 $this.bytes[$this.idx..$this.idx + $len].copy_from_slice(bytes.as_slice());
137 $this.idx_add($len);
138 Ok($len)
139 }};
140 }
141
142 pub struct WritableBytes<'a> {
143 bytes: &'a mut [u8],
144 idx: usize,
145 }
146
147 impl<'a> WritableBytes<'a> {
from(bytes: &'a mut [u8]) -> WritableBytes<'a>148 pub fn from(bytes: &'a mut [u8]) -> WritableBytes<'a> {
149 Self { bytes, idx: 0 }
150 }
151
remaining(&self) -> usize152 pub fn remaining(&self) -> usize {
153 self.bytes.len() - self.idx
154 }
155
index(&self) -> usize156 pub fn index(&self) -> usize {
157 self.idx
158 }
159
idx_add(&mut self, offset: usize)160 pub fn idx_add(&mut self, offset: usize) {
161 self.idx += offset
162 }
163
164 /// Writes an unsigned 8-bit integer at the current offset and advances
165 /// the buffer.
write_u8(&mut self, value: u8) -> Result<usize>166 pub fn write_u8(&mut self, value: u8) -> Result<usize> {
167 write_bytes!(self, value, 1)
168 }
169
write_u16(&mut self, value: u16) -> Result<usize>170 pub fn write_u16(&mut self, value: u16) -> Result<usize> {
171 write_bytes!(self, value, 2)
172 }
173
write_u32(&mut self, value: u32) -> Result<usize>174 pub fn write_u32(&mut self, value: u32) -> Result<usize> {
175 write_bytes!(self, value, 4)
176 }
177
write_u64(&mut self, value: u64) -> Result<usize>178 pub fn write_u64(&mut self, value: u64) -> Result<usize> {
179 write_bytes!(self, value, 8)
180 }
181
182 /// Writes an unsigned variable-length integer in network byte-order at the
183 /// current offset and advances the buffer.
write_varint(&mut self, value: u64) -> Result<usize>184 pub fn write_varint(&mut self, value: u64) -> Result<usize> {
185 self.write_varint_with_len(value, varint_len(value))
186 }
187
write_varint_with_len(&mut self, value: u64, len: usize) -> Result<usize>188 fn write_varint_with_len(&mut self, value: u64, len: usize) -> Result<usize> {
189 if self.remaining() < len {
190 return Err(BufferTooShort.into());
191 }
192 match len {
193 1 => self.write_u8(value as u8),
194 2 => self.write_u16(value as u16),
195 4 => self.write_u32(value as u32),
196 8 => self.write_u64(value),
197 _ => panic!("value is too large for varint"),
198 }
199 }
200 }
201
202 /// Returns how many bytes it would take to encode `v` as a variable-length
203 /// integer.
varint_len(v: u64) -> usize204 pub const fn varint_len(v: u64) -> usize {
205 match v {
206 0..=63 => 1,
207 64..=16383 => 2,
208 16384..=1_073_741_823 => 4,
209 1_073_741_824..=4_611_686_018_427_387_903 => 8,
210 _ => {
211 unreachable!()
212 }
213 }
214 }
215
216 /// Returns how long the variable-length integer is, given its first byte.
parse_varint_len(byte: u8) -> usize217 pub const fn parse_varint_len(byte: u8) -> usize {
218 let pre = byte >> 6;
219 if pre <= 3 {
220 1 << pre
221 } else {
222 unreachable!()
223 }
224 }
225
226 #[cfg(test)]
227 mod h3_octets {
228 use crate::h3::octets::{ReadableBytes, WritableBytes};
229
230 /// UT test cases for `ReadableBytes::get_varint` .
231 ///
232 /// # Brief
233 /// 1. Creates a buf with variable-length integer encoded bytes.
234 /// 2. Creates a `ReadableBytes` with the buf.
235 /// 3. Calls the get_varint method to get the value.
236 /// 4. Checks whether the result is correct.
237 #[test]
ut_readable_bytes()238 fn ut_readable_bytes() {
239 let bytes = [
240 0x3F, 0x7F, 0xFF, 0xBF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
241 0xFF,
242 ];
243 let mut readable = ReadableBytes::from(&bytes);
244 assert_eq!(readable.get_varint().unwrap(), 63);
245 assert_eq!(readable.get_varint().unwrap(), 16383);
246 assert_eq!(readable.get_varint().unwrap(), 1_073_741_823);
247 assert_eq!(readable.get_varint().unwrap(), 4_611_686_018_427_387_903);
248 assert_eq!(readable.index(), 15);
249 }
250
251 /// UT test cases for `WritableBytes::write_varint` .
252 ///
253 /// # Brief
254 /// 1. Creates a result buf with variable-length integer encoded bytes.
255 /// 2. Creates a `WritableBytes` with empty buf.
256 /// 3. Calls the write_varint method to write integers.
257 /// 4. Checks whether the result is correct.
258 #[test]
ut_writable_bytes()259 fn ut_writable_bytes() {
260 let bytes = [
261 0x3F, 0x7F, 0xFF, 0xBF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
262 0xFF,
263 ];
264 let mut buf = [0u8; 15];
265 let mut writable = WritableBytes::from(&mut buf);
266 assert_eq!(writable.write_varint(63).unwrap(), 1);
267 assert_eq!(writable.write_varint(16383).unwrap(), 2);
268 assert_eq!(writable.write_varint(1_073_741_823).unwrap(), 4);
269 assert_eq!(writable.write_varint(4_611_686_018_427_387_903).unwrap(), 8);
270 assert_eq!(writable.index(), 15);
271 assert_eq!(buf, bytes);
272 }
273 }
274