leb128_tokio/
lib.rs

1#![allow(clippy::cast_possible_truncation)]
2
3use ::core::fmt::Display;
4use ::core::future::Future;
5
6use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
7use tokio_util::bytes::{Buf as _, BytesMut};
8use tokio_util::codec::{Decoder, Encoder};
9
10/// Error returned for overflows decoding statically-sized integers
11#[derive(Debug)]
12pub struct Overflow<const N: usize>;
13
14impl Display for Overflow<8> {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        write!(f, "varint overflows an 8-bit integer")
17    }
18}
19
20impl std::error::Error for Overflow<8> {}
21
22impl Display for Overflow<16> {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        write!(f, "varint overflows a 16-bit integer")
25    }
26}
27
28impl std::error::Error for Overflow<16> {}
29
30impl Display for Overflow<32> {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "varint overflows a 32-bit integer")
33    }
34}
35
36impl std::error::Error for Overflow<32> {}
37
38impl Display for Overflow<64> {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        write!(f, "varint overflows a 64-bit integer")
41    }
42}
43
44impl std::error::Error for Overflow<64> {}
45
46impl Display for Overflow<128> {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        write!(f, "varint overflows a 128-bit integer")
49    }
50}
51
52impl std::error::Error for Overflow<128> {}
53
54/// Error returned for overflows decoding variable size integers
55#[derive(Debug)]
56pub struct OverflowVar(u8);
57
58impl Display for OverflowVar {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        write!(f, "varint overflows a {}-bit integer", self.0)
61    }
62}
63
64impl std::error::Error for OverflowVar {}
65
66fn invalid_data(err: impl Sync + Send + std::error::Error + 'static) -> std::io::Error {
67    std::io::Error::new(std::io::ErrorKind::InvalidData, err)
68}
69
70pub trait AsyncReadLeb128: AsyncRead {
71    #[cfg_attr(
72        feature = "tracing",
73        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u8"))
74    )]
75    fn read_u8_leb128(&mut self) -> impl Future<Output = std::io::Result<u8>>
76    where
77        Self: Unpin,
78    {
79        async move {
80            let mut x = 0;
81            let mut s = 0u8;
82            for _ in 0..2 {
83                let b = self.read_u8().await?;
84                if s == 7 && b > 0x01 {
85                    return Err(invalid_data(Overflow::<8>));
86                }
87                x |= (b & 0x7f) << s;
88                if b & 0x80 == 0 {
89                    return Ok(x);
90                }
91                s += 7;
92            }
93            Err(invalid_data(Overflow::<8>))
94        }
95    }
96
97    #[cfg_attr(
98        feature = "tracing",
99        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u16"))
100    )]
101    fn read_u16_leb128(&mut self) -> impl Future<Output = std::io::Result<u16>>
102    where
103        Self: Unpin,
104    {
105        async move {
106            let mut x = 0;
107            let mut s = 0u8;
108            for _ in 0..3 {
109                let b = self.read_u8().await?;
110                if s == 14 && b > 0x03 {
111                    return Err(invalid_data(Overflow::<16>));
112                }
113                x |= (u16::from(b) & 0x7f) << s;
114                if b & 0x80 == 0 {
115                    return Ok(x);
116                }
117                s += 7;
118            }
119            Err(invalid_data(Overflow::<16>))
120        }
121    }
122
123    #[cfg_attr(
124        feature = "tracing",
125        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u32"))
126    )]
127    fn read_u32_leb128(&mut self) -> impl Future<Output = std::io::Result<u32>>
128    where
129        Self: Unpin,
130    {
131        async move {
132            let mut x = 0;
133            let mut s = 0u8;
134            for _ in 0..5 {
135                let b = self.read_u8().await?;
136                if s == 28 && b > 0x0f {
137                    return Err(invalid_data(Overflow::<32>));
138                }
139                x |= (u32::from(b) & 0x7f) << s;
140                if b & 0x80 == 0 {
141                    return Ok(x);
142                }
143                s += 7;
144            }
145            Err(invalid_data(Overflow::<32>))
146        }
147    }
148
149    #[cfg_attr(
150        feature = "tracing",
151        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u64"))
152    )]
153    fn read_u64_leb128(&mut self) -> impl Future<Output = std::io::Result<u64>>
154    where
155        Self: Unpin,
156    {
157        async move {
158            let mut x = 0;
159            let mut s = 0u8;
160            for _ in 0..10 {
161                let b = self.read_u8().await?;
162                if s == 63 && b > 0x01 {
163                    return Err(invalid_data(Overflow::<64>));
164                }
165                x |= (u64::from(b) & 0x7f) << s;
166                if b & 0x80 == 0 {
167                    return Ok(x);
168                }
169                s += 7;
170            }
171            Err(invalid_data(Overflow::<64>))
172        }
173    }
174
175    #[cfg_attr(
176        feature = "tracing",
177        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u128"))
178    )]
179    fn read_u128_leb128(&mut self) -> impl Future<Output = std::io::Result<u128>>
180    where
181        Self: Unpin,
182    {
183        async move {
184            let mut x = 0;
185            let mut s = 0u8;
186            for _ in 0..19 {
187                let b = self.read_u8().await?;
188                if s == 126 && b > 0x03 {
189                    return Err(invalid_data(Overflow::<128>));
190                }
191                x |= (u128::from(b) & 0x7f) << s;
192                if b & 0x80 == 0 {
193                    return Ok(x);
194                }
195                s += 7;
196            }
197            Err(invalid_data(Overflow::<128>))
198        }
199    }
200
201    #[cfg_attr(
202        feature = "tracing",
203        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "uvarint", n))
204    )]
205    fn read_var_u8_leb128(&mut self, n: u8) -> impl Future<Output = std::io::Result<u8>>
206    where
207        Self: Unpin,
208    {
209        async move {
210            if n == 8 {
211                return self.read_u8_leb128().await;
212            }
213            debug_assert!(n <= 8);
214            let max = (n / 7) + 1;
215            let mut x = 0u8;
216            let mut s = 0u8;
217            for _ in 0..max {
218                let b = self.read_u8().await?;
219                if s == (n / 7) * 7 && b > n % 7 {
220                    return Err(invalid_data(OverflowVar(n)));
221                }
222                x |= (b & 0x7f) << s;
223                if b & 0x80 == 0 {
224                    return Ok(x);
225                }
226                s += 7;
227            }
228            Err(invalid_data(OverflowVar(n)))
229        }
230    }
231
232    #[cfg_attr(
233        feature = "tracing",
234        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "uvarint", n))
235    )]
236    fn read_var_u16_leb128(&mut self, n: u8) -> impl Future<Output = std::io::Result<u16>>
237    where
238        Self: Unpin,
239    {
240        async move {
241            match n {
242                8 => return self.read_u8_leb128().await.map(Into::into),
243                16 => return self.read_u16_leb128().await,
244                _ => {}
245            }
246            debug_assert!(n <= 16);
247            let max = (n / 7) + 1;
248            let mut x = 0;
249            let mut s = 0u8;
250            for _ in 0..max {
251                let b = self.read_u8().await?;
252                if s == (n / 7) * 7 && b > n % 7 {
253                    return Err(invalid_data(OverflowVar(n)));
254                }
255                x |= (u16::from(b) & 0x7f) << s;
256                if b & 0x80 == 0 {
257                    return Ok(x);
258                }
259                s += 7;
260            }
261            Err(invalid_data(OverflowVar(n)))
262        }
263    }
264
265    #[cfg_attr(
266        feature = "tracing",
267        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "uvarint", n))
268    )]
269    fn read_var_u32_leb128(&mut self, n: u8) -> impl Future<Output = std::io::Result<u32>>
270    where
271        Self: Unpin,
272    {
273        async move {
274            match n {
275                8 => return self.read_u8_leb128().await.map(Into::into),
276                16 => return self.read_u16_leb128().await.map(Into::into),
277                32 => return self.read_u32_leb128().await,
278                _ => {}
279            }
280            debug_assert!(n <= 32);
281            let max = (n / 7) + 1;
282            let mut x = 0;
283            let mut s = 0u8;
284            for _ in 0..max {
285                let b = self.read_u8().await?;
286                if s == (n / 7) * 7 && b > n % 7 {
287                    return Err(invalid_data(OverflowVar(n)));
288                }
289                x |= (u32::from(b) & 0x7f) << s;
290                if b & 0x80 == 0 {
291                    return Ok(x);
292                }
293                s += 7;
294            }
295            Err(invalid_data(OverflowVar(n)))
296        }
297    }
298
299    #[cfg_attr(
300        feature = "tracing",
301        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "uvarint", n))
302    )]
303    fn read_var_u64_leb128(&mut self, n: u8) -> impl Future<Output = std::io::Result<u64>>
304    where
305        Self: Unpin,
306    {
307        async move {
308            match n {
309                8 => return self.read_u8_leb128().await.map(Into::into),
310                16 => return self.read_u16_leb128().await.map(Into::into),
311                32 => return self.read_u32_leb128().await.map(Into::into),
312                64 => return self.read_u64_leb128().await,
313                _ => {}
314            }
315            debug_assert!(n <= 64);
316            let max = (n / 7) + 1;
317            let mut x = 0;
318            let mut s = 0u8;
319            for _ in 0..max {
320                let b = self.read_u8().await?;
321                if s == (n / 7) * 7 && b > n % 7 {
322                    return Err(invalid_data(OverflowVar(n)));
323                }
324                x |= (u64::from(b) & 0x7f) << s;
325                if b & 0x80 == 0 {
326                    return Ok(x);
327                }
328                s += 7;
329            }
330            Err(invalid_data(OverflowVar(n)))
331        }
332    }
333
334    #[cfg_attr(
335        feature = "tracing",
336        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "uvarint", n))
337    )]
338    fn read_var_u128_leb128(&mut self, n: u8) -> impl Future<Output = std::io::Result<u128>>
339    where
340        Self: Unpin,
341    {
342        async move {
343            match n {
344                8 => return self.read_u8_leb128().await.map(Into::into),
345                16 => return self.read_u16_leb128().await.map(Into::into),
346                32 => return self.read_u32_leb128().await.map(Into::into),
347                64 => return self.read_u64_leb128().await.map(Into::into),
348                128 => return self.read_u128_leb128().await,
349                _ => {}
350            }
351            debug_assert!(n <= 128);
352            let max = (n / 7) + 1;
353            let mut x = 0;
354            let mut s = 0u8;
355            for _ in 0..max {
356                let b = self.read_u8().await?;
357                if s == (n / 7) * 7 && b > n % 7 {
358                    return Err(invalid_data(OverflowVar(n)));
359                }
360                x |= (u128::from(b) & 0x7f) << s;
361                if b & 0x80 == 0 {
362                    return Ok(x);
363                }
364                s += 7;
365            }
366            Err(invalid_data(OverflowVar(n)))
367        }
368    }
369
370    #[cfg_attr(
371        feature = "tracing",
372        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i8"))
373    )]
374    fn read_i8_leb128(&mut self) -> impl Future<Output = std::io::Result<i8>>
375    where
376        Self: Unpin,
377    {
378        async move {
379            let mut x = 0;
380            let mut s = 0u8;
381            for _ in 0..2 {
382                let b = self.read_u8().await?;
383                if s == 7 && b > 0x01 {
384                    return Err(invalid_data(Overflow::<8>));
385                }
386                x |= ((b as i8) & 0x7f) << s;
387                s += 7;
388                if b & 0x80 == 0 {
389                    if s != 14 && b & 0x40 != 0 {
390                        return Ok(x | !0 << s);
391                    } else {
392                        return Ok(x);
393                    }
394                }
395            }
396            Err(invalid_data(Overflow::<8>))
397        }
398    }
399
400    #[cfg_attr(
401        feature = "tracing",
402        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i16"))
403    )]
404    fn read_i16_leb128(&mut self) -> impl Future<Output = std::io::Result<i16>>
405    where
406        Self: Unpin,
407    {
408        async move {
409            let mut x = 0;
410            let mut s = 0u8;
411            for _ in 0..3 {
412                let b = self.read_u8().await?;
413                if s == 14 && b > 0x03 {
414                    return Err(invalid_data(Overflow::<16>));
415                }
416                x |= (i16::from(b) & 0x7f) << s;
417                s += 7;
418                if b & 0x80 == 0 {
419                    if s != 21 && b & 0x40 != 0 {
420                        return Ok(x | !0 << s);
421                    } else {
422                        return Ok(x);
423                    }
424                }
425            }
426            Err(invalid_data(Overflow::<16>))
427        }
428    }
429
430    #[cfg_attr(
431        feature = "tracing",
432        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i32"))
433    )]
434    fn read_i32_leb128(&mut self) -> impl Future<Output = std::io::Result<i32>>
435    where
436        Self: Unpin,
437    {
438        async move {
439            let mut x = 0;
440            let mut s = 0u8;
441            for _ in 0..5 {
442                let b = self.read_u8().await?;
443                if s == 28 && b > 0x0f {
444                    return Err(invalid_data(Overflow::<32>));
445                }
446                x |= (i32::from(b) & 0x7f) << s;
447                s += 7;
448                if b & 0x80 == 0 {
449                    if s != 35 && b & 0x40 != 0 {
450                        return Ok(x | !0 << s);
451                    } else {
452                        return Ok(x);
453                    }
454                }
455            }
456            Err(invalid_data(Overflow::<32>))
457        }
458    }
459
460    #[cfg_attr(
461        feature = "tracing",
462        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i64"))
463    )]
464    fn read_i64_leb128(&mut self) -> impl Future<Output = std::io::Result<i64>>
465    where
466        Self: Unpin,
467    {
468        async move {
469            let mut x = 0;
470            let mut s = 0u8;
471            for _ in 0..10 {
472                let b = self.read_u8().await?;
473                if s == 63 && b > 0x01 {
474                    return Err(invalid_data(Overflow::<64>));
475                }
476                x |= (i64::from(b) & 0x7f) << s;
477                s += 7;
478                if b & 0x80 == 0 {
479                    if s != 70 && b & 0x40 != 0 {
480                        return Ok(x | !0 << s);
481                    } else {
482                        return Ok(x);
483                    }
484                }
485            }
486            Err(invalid_data(Overflow::<64>))
487        }
488    }
489
490    #[cfg_attr(
491        feature = "tracing",
492        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i128"))
493    )]
494    fn read_i128_leb128(&mut self) -> impl Future<Output = std::io::Result<i128>>
495    where
496        Self: Unpin,
497    {
498        async move {
499            let mut x = 0;
500            let mut s = 0u8;
501            for _ in 0..19 {
502                let b = self.read_u8().await?;
503                if s == 126 && b > 0x03 {
504                    return Err(invalid_data(Overflow::<128>));
505                }
506                x |= (i128::from(b) & 0x7f) << s;
507                s += 7;
508                if b & 0x80 == 0 {
509                    if s != 133 && b & 0x40 != 0 {
510                        return Ok(x | !0 << s);
511                    } else {
512                        return Ok(x);
513                    }
514                }
515            }
516            Err(invalid_data(Overflow::<128>))
517        }
518    }
519}
520
521impl<T: AsyncRead> AsyncReadLeb128 for T {}
522
523pub fn put_u8_leb128(buf: &mut [u8; 2], mut x: u8) -> &mut [u8] {
524    let mut i = 0;
525    while x >= 0x80 {
526        buf[i] = x | 0x80;
527        x >>= 7;
528        i += 1;
529    }
530    buf[i] = x;
531    &mut buf[..=i]
532}
533
534pub fn put_u16_leb128(buf: &mut [u8; 3], mut x: u16) -> &mut [u8] {
535    let mut i = 0;
536    while x >= 0x80 {
537        buf[i] = (x as u8) | 0x80;
538        x >>= 7;
539        i += 1;
540    }
541    buf[i] = x as u8;
542    &mut buf[..=i]
543}
544
545pub fn put_u32_leb128(buf: &mut [u8; 5], mut x: u32) -> &mut [u8] {
546    let mut i = 0;
547    while x >= 0x80 {
548        buf[i] = (x as u8) | 0x80;
549        x >>= 7;
550        i += 1;
551    }
552    buf[i] = x as u8;
553    &mut buf[..=i]
554}
555
556pub fn put_u64_leb128(buf: &mut [u8; 10], mut x: u64) -> &mut [u8] {
557    let mut i = 0;
558    while x >= 0x80 {
559        buf[i] = (x as u8) | 0x80;
560        x >>= 7;
561        i += 1;
562    }
563    buf[i] = x as u8;
564    &mut buf[..=i]
565}
566
567pub fn put_u128_leb128(buf: &mut [u8; 19], mut x: u128) -> &mut [u8] {
568    let mut i = 0;
569    while x >= 0x80 {
570        buf[i] = (x as u8) | 0x80;
571        x >>= 7;
572        i += 1;
573    }
574    buf[i] = x as u8;
575    &mut buf[..=i]
576}
577
578pub fn put_i8_leb128(buf: &mut [u8; 2], mut x: i8) -> &mut [u8] {
579    let mut i = 0;
580    loop {
581        let b = x as u8;
582        x >>= 6;
583        if x == 0 || x == !0 {
584            buf[i] = b & 0x7f;
585            return &mut buf[..=i];
586        } else {
587            buf[i] = b | 0x80;
588            x >>= 1;
589        }
590        i += 1;
591    }
592}
593
594pub fn put_i16_leb128(buf: &mut [u8; 3], mut x: i16) -> &mut [u8] {
595    let mut i = 0;
596    loop {
597        let b = x as u8;
598        x >>= 6;
599        if x == 0 || x == !0 {
600            buf[i] = b & 0x7f;
601            return &mut buf[..=i];
602        } else {
603            buf[i] = b | 0x80;
604            x >>= 1;
605        }
606        i += 1;
607    }
608}
609
610pub fn put_i32_leb128(buf: &mut [u8; 5], mut x: i32) -> &mut [u8] {
611    let mut i = 0;
612    loop {
613        let b = x as u8;
614        x >>= 6;
615        if x == 0 || x == !0 {
616            buf[i] = b & 0x7f;
617            return &mut buf[..=i];
618        } else {
619            buf[i] = b | 0x80;
620            x >>= 1;
621        }
622        i += 1;
623    }
624}
625
626pub fn put_i64_leb128(buf: &mut [u8; 10], mut x: i64) -> &mut [u8] {
627    let mut i = 0;
628    loop {
629        let b = x as u8;
630        x >>= 6;
631        if x == 0 || x == !0 {
632            buf[i] = b & 0x7f;
633            return &mut buf[..=i];
634        } else {
635            buf[i] = b | 0x80;
636            x >>= 1;
637        }
638        i += 1;
639    }
640}
641
642pub fn put_i128_leb128(buf: &mut [u8; 19], mut x: i128) -> &mut [u8] {
643    let mut i = 0;
644    loop {
645        let b = x as u8;
646        x >>= 6;
647        if x == 0 || x == !0 {
648            buf[i] = b & 0x7f;
649            return &mut buf[..=i];
650        } else {
651            buf[i] = b | 0x80;
652            x >>= 1;
653        }
654        i += 1;
655    }
656}
657
658pub trait AsyncWriteLeb128: AsyncWrite {
659    #[cfg_attr(
660        feature = "tracing",
661        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u8"))
662    )]
663    fn write_u8_leb128(&mut self, x: u8) -> impl Future<Output = std::io::Result<()>>
664    where
665        Self: Unpin,
666    {
667        async move {
668            self.write_all(put_u8_leb128(&mut Default::default(), x))
669                .await
670        }
671    }
672
673    #[cfg_attr(
674        feature = "tracing",
675        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u16"))
676    )]
677    fn write_u16_leb128(&mut self, x: u16) -> impl Future<Output = std::io::Result<()>>
678    where
679        Self: Unpin,
680    {
681        async move {
682            self.write_all(put_u16_leb128(&mut Default::default(), x))
683                .await
684        }
685    }
686
687    #[cfg_attr(
688        feature = "tracing",
689        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u32"))
690    )]
691    fn write_u32_leb128(&mut self, x: u32) -> impl Future<Output = std::io::Result<()>>
692    where
693        Self: Unpin,
694    {
695        async move {
696            self.write_all(put_u32_leb128(&mut Default::default(), x))
697                .await
698        }
699    }
700
701    #[cfg_attr(
702        feature = "tracing",
703        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u64"))
704    )]
705    fn write_u64_leb128(&mut self, x: u64) -> impl Future<Output = std::io::Result<()>>
706    where
707        Self: Unpin,
708    {
709        async move {
710            self.write_all(put_u64_leb128(&mut Default::default(), x))
711                .await
712        }
713    }
714
715    #[cfg_attr(
716        feature = "tracing",
717        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "u128"))
718    )]
719    fn write_u128_leb128(&mut self, x: u128) -> impl Future<Output = std::io::Result<()>>
720    where
721        Self: Unpin,
722    {
723        async move {
724            self.write_all(put_u128_leb128(&mut Default::default(), x))
725                .await
726        }
727    }
728
729    #[cfg_attr(
730        feature = "tracing",
731        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i8"))
732    )]
733    fn write_i8_leb128(&mut self, x: i8) -> impl Future<Output = std::io::Result<()>>
734    where
735        Self: Unpin,
736    {
737        async move {
738            self.write_all(put_i8_leb128(&mut Default::default(), x))
739                .await
740        }
741    }
742
743    #[cfg_attr(
744        feature = "tracing",
745        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i16"))
746    )]
747    fn write_i16_leb128(&mut self, x: i16) -> impl Future<Output = std::io::Result<()>>
748    where
749        Self: Unpin,
750    {
751        async move {
752            self.write_all(put_i16_leb128(&mut Default::default(), x))
753                .await
754        }
755    }
756
757    #[cfg_attr(
758        feature = "tracing",
759        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i32"))
760    )]
761    fn write_i32_leb128(&mut self, x: i32) -> impl Future<Output = std::io::Result<()>>
762    where
763        Self: Unpin,
764    {
765        async move {
766            self.write_all(put_i32_leb128(&mut Default::default(), x))
767                .await
768        }
769    }
770
771    #[cfg_attr(
772        feature = "tracing",
773        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i64"))
774    )]
775    fn write_i64_leb128(&mut self, x: i64) -> impl Future<Output = std::io::Result<()>>
776    where
777        Self: Unpin,
778    {
779        async move {
780            self.write_all(put_i64_leb128(&mut Default::default(), x))
781                .await
782        }
783    }
784
785    #[cfg_attr(
786        feature = "tracing",
787        tracing::instrument(level = "trace", ret, skip_all, fields(ty = "i128"))
788    )]
789    fn write_i128_leb128(&mut self, x: i128) -> impl Future<Output = std::io::Result<()>>
790    where
791        Self: Unpin,
792    {
793        async move {
794            self.write_all(put_i128_leb128(&mut Default::default(), x))
795                .await
796        }
797    }
798}
799
800impl<T: AsyncWrite> AsyncWriteLeb128 for T {}
801
802pub struct Leb128DecoderU8;
803
804impl Decoder for Leb128DecoderU8 {
805    type Item = u8;
806    type Error = std::io::Error;
807
808    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
809        let mut x = 0;
810        let mut s = 0u8;
811        for i in 0..2 {
812            let Some(b) = src.get(i) else {
813                src.reserve(1);
814                return Ok(None);
815            };
816            if s == 7 && *b > 0x01 {
817                return Err(invalid_data(Overflow::<8>));
818            }
819            x |= (b & 0x7f) << s;
820            if b & 0x80 == 0 {
821                src.advance(i + 1);
822                return Ok(Some(x));
823            }
824            s += 7;
825        }
826        Err(invalid_data(Overflow::<8>))
827    }
828}
829
830pub struct Leb128DecoderU16;
831
832impl Decoder for Leb128DecoderU16 {
833    type Item = u16;
834    type Error = std::io::Error;
835
836    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
837        let mut x = 0;
838        let mut s = 0u8;
839        for i in 0..3 {
840            let Some(b) = src.get(i) else {
841                src.reserve(1);
842                return Ok(None);
843            };
844            if s == 14 && *b > 0x03 {
845                return Err(invalid_data(Overflow::<16>));
846            }
847            x |= (u16::from(*b) & 0x7f) << s;
848            if b & 0x80 == 0 {
849                src.advance(i + 1);
850                return Ok(Some(x));
851            }
852            s += 7;
853        }
854        Err(invalid_data(Overflow::<16>))
855    }
856}
857
858pub struct Leb128DecoderU32;
859
860impl Decoder for Leb128DecoderU32 {
861    type Item = u32;
862    type Error = std::io::Error;
863
864    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
865        let mut x = 0;
866        let mut s = 0u8;
867        for i in 0..5 {
868            let Some(b) = src.get(i) else {
869                src.reserve(1);
870                return Ok(None);
871            };
872            if s == 28 && *b > 0x0f {
873                return Err(invalid_data(Overflow::<32>));
874            }
875            x |= (u32::from(*b) & 0x7f) << s;
876            if b & 0x80 == 0 {
877                src.advance(i + 1);
878                return Ok(Some(x));
879            }
880            s += 7;
881        }
882        Err(invalid_data(Overflow::<32>))
883    }
884}
885
886pub struct Leb128DecoderU64;
887
888impl Decoder for Leb128DecoderU64 {
889    type Item = u64;
890    type Error = std::io::Error;
891
892    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
893        let mut x = 0;
894        let mut s = 0u8;
895        for i in 0..10 {
896            let Some(b) = src.get(i) else {
897                src.reserve(1);
898                return Ok(None);
899            };
900            if s == 63 && *b > 0x01 {
901                return Err(invalid_data(Overflow::<64>));
902            }
903            x |= (u64::from(*b) & 0x7f) << s;
904            if b & 0x80 == 0 {
905                src.advance(i + 1);
906                return Ok(Some(x));
907            }
908            s += 7;
909        }
910        Err(invalid_data(Overflow::<64>))
911    }
912}
913
914pub struct Leb128DecoderU128;
915
916impl Decoder for Leb128DecoderU128 {
917    type Item = u128;
918    type Error = std::io::Error;
919
920    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
921        let mut x = 0;
922        let mut s = 0u8;
923        for i in 0..19 {
924            let Some(b) = src.get(i) else {
925                src.reserve(1);
926                return Ok(None);
927            };
928            if s == 126 && *b > 0x03 {
929                return Err(invalid_data(Overflow::<128>));
930            }
931            x |= (u128::from(*b) & 0x7f) << s;
932            if b & 0x80 == 0 {
933                src.advance(i + 1);
934                return Ok(Some(x));
935            }
936            s += 7;
937        }
938        Err(invalid_data(Overflow::<128>))
939    }
940}
941
942pub struct Leb128DecoderI8;
943
944impl Decoder for Leb128DecoderI8 {
945    type Item = i8;
946    type Error = std::io::Error;
947
948    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
949        let mut x = 0;
950        let mut s = 0u8;
951        for i in 0..2 {
952            let Some(b) = src.get(i).copied() else {
953                src.reserve(1);
954                return Ok(None);
955            };
956            if s == 7 && b > 0x01 {
957                return Err(invalid_data(Overflow::<8>));
958            }
959            x |= ((b as i8) & 0x7f) << s;
960            s += 7;
961            if b & 0x80 == 0 {
962                src.advance(i + 1);
963                if s != 14 && b & 0x40 != 0 {
964                    return Ok(Some(x | !0 << s));
965                } else {
966                    return Ok(Some(x));
967                }
968            }
969        }
970        Err(invalid_data(Overflow::<8>))
971    }
972}
973
974pub struct Leb128DecoderI16;
975
976impl Decoder for Leb128DecoderI16 {
977    type Item = i16;
978    type Error = std::io::Error;
979
980    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
981        let mut x = 0;
982        let mut s = 0u8;
983        for i in 0..3 {
984            let Some(b) = src.get(i).copied() else {
985                src.reserve(1);
986                return Ok(None);
987            };
988            if s == 14 && b > 0x03 {
989                return Err(invalid_data(Overflow::<16>));
990            }
991            x |= (i16::from(b) & 0x7f) << s;
992            s += 7;
993            if b & 0x80 == 0 {
994                src.advance(i + 1);
995                if s != 21 && b & 0x40 != 0 {
996                    return Ok(Some(x | !0 << s));
997                } else {
998                    return Ok(Some(x));
999                }
1000            }
1001        }
1002        Err(invalid_data(Overflow::<16>))
1003    }
1004}
1005
1006pub struct Leb128DecoderI32;
1007
1008impl Decoder for Leb128DecoderI32 {
1009    type Item = i32;
1010    type Error = std::io::Error;
1011
1012    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
1013        let mut x = 0;
1014        let mut s = 0u8;
1015        for i in 0..5 {
1016            let Some(b) = src.get(i).copied() else {
1017                src.reserve(1);
1018                return Ok(None);
1019            };
1020            if s == 28 && b > 0x0f {
1021                return Err(invalid_data(Overflow::<32>));
1022            }
1023            x |= (i32::from(b) & 0x7f) << s;
1024            s += 7;
1025            if b & 0x80 == 0 {
1026                src.advance(i + 1);
1027                if s != 35 && b & 0x40 != 0 {
1028                    return Ok(Some(x | !0 << s));
1029                } else {
1030                    return Ok(Some(x));
1031                }
1032            }
1033        }
1034        Err(invalid_data(Overflow::<32>))
1035    }
1036}
1037
1038pub struct Leb128DecoderI64;
1039
1040impl Decoder for Leb128DecoderI64 {
1041    type Item = i64;
1042    type Error = std::io::Error;
1043
1044    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
1045        let mut x = 0;
1046        let mut s = 0u8;
1047        for i in 0..10 {
1048            let Some(b) = src.get(i).copied() else {
1049                src.reserve(1);
1050                return Ok(None);
1051            };
1052            if s == 63 && b > 0x01 {
1053                return Err(invalid_data(Overflow::<64>));
1054            }
1055            x |= (i64::from(b) & 0x7f) << s;
1056            s += 7;
1057            if b & 0x80 == 0 {
1058                src.advance(i + 1);
1059                if s != 70 && b & 0x40 != 0 {
1060                    return Ok(Some(x | !0 << s));
1061                } else {
1062                    return Ok(Some(x));
1063                }
1064            }
1065        }
1066        Err(invalid_data(Overflow::<64>))
1067    }
1068}
1069
1070pub struct Leb128DecoderI128;
1071
1072impl Decoder for Leb128DecoderI128 {
1073    type Item = i128;
1074    type Error = std::io::Error;
1075
1076    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
1077        let mut x = 0;
1078        let mut s = 0u8;
1079        for i in 0..19 {
1080            let Some(b) = src.get(i).copied() else {
1081                src.reserve(1);
1082                return Ok(None);
1083            };
1084            if s == 126 && b > 0x03 {
1085                return Err(invalid_data(Overflow::<128>));
1086            }
1087            x |= (i128::from(b) & 0x7f) << s;
1088            s += 7;
1089            if b & 0x80 == 0 {
1090                src.advance(i + 1);
1091                if s != 133 && b & 0x40 != 0 {
1092                    return Ok(Some(x | !0 << s));
1093                } else {
1094                    return Ok(Some(x));
1095                }
1096            }
1097        }
1098        Err(invalid_data(Overflow::<128>))
1099    }
1100}
1101
1102pub struct Leb128Encoder;
1103
1104macro_rules! impl_encode {
1105    ($t:ty, $f:ident) => {
1106        impl Encoder<$t> for Leb128Encoder {
1107            type Error = std::io::Error;
1108
1109            fn encode(&mut self, item: $t, dst: &mut BytesMut) -> Result<(), Self::Error> {
1110                dst.extend_from_slice($f(&mut Default::default(), item));
1111                Ok(())
1112            }
1113        }
1114
1115        impl Encoder<&$t> for Leb128Encoder {
1116            type Error = std::io::Error;
1117
1118            fn encode(&mut self, item: &$t, dst: &mut BytesMut) -> Result<(), Self::Error> {
1119                self.encode(*item, dst)
1120            }
1121        }
1122
1123        impl Encoder<&&$t> for Leb128Encoder {
1124            type Error = std::io::Error;
1125
1126            fn encode(&mut self, item: &&$t, dst: &mut BytesMut) -> Result<(), Self::Error> {
1127                self.encode(*item, dst)
1128            }
1129        }
1130    };
1131}
1132
1133impl_encode!(u8,   put_u8_leb128);
1134impl_encode!(u16,  put_u16_leb128);
1135impl_encode!(u32,  put_u32_leb128);
1136impl_encode!(u64,  put_u64_leb128);
1137impl_encode!(u128, put_u128_leb128);
1138impl_encode!(i8,   put_i8_leb128);
1139impl_encode!(i16,  put_i16_leb128);
1140impl_encode!(i32,  put_i32_leb128);
1141impl_encode!(i64,  put_i64_leb128);
1142impl_encode!(i128, put_i128_leb128);
1143
1144#[cfg(test)]
1145mod tests {
1146    use super::*;
1147
1148    #[tokio::test]
1149    async fn unsigned_leb128() {
1150        const ENCODED: [u8; 3] = [0xe5, 0x8e, 0x26];
1151
1152        let v = ENCODED
1153            .as_slice()
1154            .read_u32_leb128()
1155            .await
1156            .expect("failed to read u32");
1157        assert_eq!(v, 624_485);
1158
1159        let v = Leb128DecoderU32
1160            .decode(&mut ENCODED.as_slice().into())
1161            .expect("failed to decode u32");
1162        assert_eq!(v, Some(624_485));
1163
1164        let v = ENCODED
1165            .as_slice()
1166            .read_u64_leb128()
1167            .await
1168            .expect("failed to read u64");
1169        assert_eq!(v, 624_485);
1170
1171        let v = Leb128DecoderU64
1172            .decode(&mut ENCODED.as_slice().into())
1173            .expect("failed to decode u64");
1174        assert_eq!(v, Some(624_485));
1175
1176        let mut buf = vec![];
1177        buf.write_u32_leb128(624_485)
1178            .await
1179            .expect("failed to write u32");
1180        assert_eq!(buf, ENCODED);
1181
1182        let mut buf = vec![];
1183        buf.write_u64_leb128(624_485)
1184            .await
1185            .expect("failed to write u64");
1186        assert_eq!(buf, ENCODED);
1187
1188        let v = [0xff, 0x01]
1189            .as_slice()
1190            .read_u8_leb128()
1191            .await
1192            .expect("failed to read u8");
1193        assert_eq!(v, u8::MAX);
1194
1195        let v = [0x7f]
1196            .as_slice()
1197            .read_u8_leb128()
1198            .await
1199            .expect("failed to read u8");
1200        assert_eq!(v, 0x7f);
1201
1202        let v = [0xff, 0x00]
1203            .as_slice()
1204            .read_u8_leb128()
1205            .await
1206            .expect("failed to read u8");
1207        assert_eq!(v, 0x7f);
1208
1209        [0xff, 0x02]
1210            .as_slice()
1211            .read_u8_leb128()
1212            .await
1213            .expect_err("u8 read should have failed, since it encoded 9 bits");
1214
1215        let v = [0xff, 0xff, 0x01]
1216            .as_slice()
1217            .read_u16_leb128()
1218            .await
1219            .expect("failed to read u16");
1220        assert_eq!(v, 0x7fff);
1221
1222        let v = [0xff, 0xff, 0x02]
1223            .as_slice()
1224            .read_u16_leb128()
1225            .await
1226            .expect("failed to read u16");
1227        assert_eq!(v, 0xbfff);
1228
1229        let v = [0xff, 0xff, 0x03]
1230            .as_slice()
1231            .read_u16_leb128()
1232            .await
1233            .expect("failed to read u16");
1234        assert_eq!(v, u16::MAX);
1235
1236        [0xff, 0xff, 0x04]
1237            .as_slice()
1238            .read_u16_leb128()
1239            .await
1240            .expect_err("u16 read should have failed, since it encoded 17 bits");
1241
1242        let v = [0xff, 0xff, 0xff, 0xff, 0x0f]
1243            .as_slice()
1244            .read_u32_leb128()
1245            .await
1246            .expect("failed to read u32");
1247        assert_eq!(v, u32::MAX);
1248
1249        [0xff, 0xff, 0xff, 0xff, 0x10]
1250            .as_slice()
1251            .read_u32_leb128()
1252            .await
1253            .expect_err("u32 read should have failed, since it encoded 33 bits");
1254
1255        let v = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01]
1256            .as_slice()
1257            .read_u64_leb128()
1258            .await
1259            .expect("failed to read u64");
1260        assert_eq!(v, u64::MAX);
1261
1262        [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02]
1263            .as_slice()
1264            .read_u64_leb128()
1265            .await
1266            .expect_err("u64 read should have failed, since it encoded 65 bits");
1267
1268        let v = [
1269            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
1270            0xff, 0xff, 0xff, 0xff, 0x03,
1271        ]
1272        .as_slice()
1273        .read_u128_leb128()
1274        .await
1275        .expect("failed to read u128");
1276        assert_eq!(v, u128::MAX);
1277
1278        [
1279            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
1280            0xff, 0xff, 0xff, 0xff, 0x04,
1281        ]
1282        .as_slice()
1283        .read_u128_leb128()
1284        .await
1285        .expect_err("u128 read should have failed, since it encoded 129 bits");
1286    }
1287
1288    #[tokio::test]
1289    async fn signed_leb128() {
1290        const ENCODED: [u8; 3] = [0xc0, 0xbb, 0x78];
1291
1292        let v = ENCODED
1293            .as_slice()
1294            .read_i32_leb128()
1295            .await
1296            .expect("failed to read i32");
1297        assert_eq!(v, -123_456);
1298
1299        let v = ENCODED
1300            .as_slice()
1301            .read_i64_leb128()
1302            .await
1303            .expect("failed to read i64");
1304        assert_eq!(v, -123_456);
1305
1306        let mut buf = vec![];
1307        buf.write_i32_leb128(-123_456)
1308            .await
1309            .expect("failed to write i32");
1310        assert_eq!(buf, ENCODED);
1311
1312        let mut buf = vec![];
1313        buf.write_i64_leb128(-123_456)
1314            .await
1315            .expect("failed to write i64");
1316        assert_eq!(buf, ENCODED);
1317
1318        let mut buf = vec![];
1319        buf.write_i64_leb128(-1).await.expect("failed to write i64");
1320        assert_eq!(buf, [0x7f]);
1321
1322        let mut buf = vec![];
1323        buf.write_i64_leb128(-2).await.expect("failed to write i64");
1324        assert_eq!(buf, [0x7e]);
1325
1326        let v = [0xff, 0x01]
1327            .as_slice()
1328            .read_i8_leb128()
1329            .await
1330            .expect("failed to read i8");
1331        assert_eq!(v, -1);
1332
1333        let v = [0x7f]
1334            .as_slice()
1335            .read_i8_leb128()
1336            .await
1337            .expect("failed to read i8");
1338        assert_eq!(v, -1);
1339
1340        let v = [0xff, 0x00]
1341            .as_slice()
1342            .read_i8_leb128()
1343            .await
1344            .expect("failed to read i8");
1345        assert_eq!(v, 0x7f);
1346
1347        [0xff, 0x02]
1348            .as_slice()
1349            .read_i8_leb128()
1350            .await
1351            .expect_err("i8 read should have failed, since it encoded 9 bits");
1352
1353        let v = [0xff, 0xff, 0x01]
1354            .as_slice()
1355            .read_i16_leb128()
1356            .await
1357            .expect("failed to read i16");
1358        assert_eq!(v, 0x7fff);
1359
1360        let v = [0xff, 0xff, 0x02]
1361            .as_slice()
1362            .read_i16_leb128()
1363            .await
1364            .expect("failed to read i16");
1365        assert_eq!(v, -0x4001);
1366
1367        let v = [0xff, 0xff, 0x03]
1368            .as_slice()
1369            .read_i16_leb128()
1370            .await
1371            .expect("failed to read i16");
1372        assert_eq!(v, -1);
1373
1374        [0xff, 0xff, 0x04]
1375            .as_slice()
1376            .read_i16_leb128()
1377            .await
1378            .expect_err("i16 read should have failed, since it encoded 17 bits");
1379
1380        let v = [0x7f]
1381            .as_slice()
1382            .read_i32_leb128()
1383            .await
1384            .expect("failed to read i32");
1385        assert_eq!(v, -1);
1386
1387        let v = [0x7e]
1388            .as_slice()
1389            .read_i32_leb128()
1390            .await
1391            .expect("failed to read i32");
1392        assert_eq!(v, -2);
1393
1394        let v = [0xff, 0xff, 0xff, 0xff, 0x0f]
1395            .as_slice()
1396            .read_i32_leb128()
1397            .await
1398            .expect("failed to read i32");
1399        assert_eq!(v, -1);
1400
1401        let v = [0xfe, 0xff, 0xff, 0xff, 0x0f]
1402            .as_slice()
1403            .read_i32_leb128()
1404            .await
1405            .expect("failed to read i32");
1406        assert_eq!(v, -2);
1407
1408        [0xff, 0xff, 0xff, 0xff, 0x10]
1409            .as_slice()
1410            .read_i32_leb128()
1411            .await
1412            .expect_err("i32 read should have failed, since it encoded 33 bits");
1413
1414        let v = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01]
1415            .as_slice()
1416            .read_i64_leb128()
1417            .await
1418            .expect("failed to read i64");
1419        assert_eq!(v, -1);
1420
1421        [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02]
1422            .as_slice()
1423            .read_i64_leb128()
1424            .await
1425            .expect_err("i64 read should have failed, since it encoded 65 bits");
1426
1427        let v = [
1428            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
1429            0xff, 0xff, 0xff, 0xff, 0x03,
1430        ]
1431        .as_slice()
1432        .read_i128_leb128()
1433        .await
1434        .expect("failed to read i128");
1435        assert_eq!(v, -1);
1436
1437        [
1438            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
1439            0xff, 0xff, 0xff, 0xff, 0x04,
1440        ]
1441        .as_slice()
1442        .read_i128_leb128()
1443        .await
1444        .expect_err("i128 read should have failed, since it encoded 129 bits");
1445    }
1446
1447    #[tokio::test]
1448    async fn unsigned() {
1449        let v = [0x01u8]
1450            .as_slice()
1451            .read_var_u8_leb128(2)
1452            .await
1453            .expect("failed to read u2");
1454        assert_eq!(v, 1);
1455
1456        let v = [0x02]
1457            .as_slice()
1458            .read_var_u8_leb128(2)
1459            .await
1460            .expect("failed to read u2");
1461        assert_eq!(v, 2);
1462
1463        [0b100]
1464            .as_slice()
1465            .read_var_u8_leb128(2)
1466            .await
1467            .expect_err("u2 read should have failed, since it encoded 3 bits");
1468
1469        [0x80, 0x80, 0x01]
1470            .as_slice()
1471            .read_var_u16_leb128(9)
1472            .await
1473            .expect_err("u9 read should have failed, since it used over 9 bits");
1474
1475        let v = [0x80, 0x80, 0x80, 0x80, 0x80, 0x01]
1476            .as_slice()
1477            .read_var_u64_leb128(64)
1478            .await
1479            .expect("failed to read u64");
1480        assert_eq!(v, 0b1000_0000_0000_0000_0000_0000_0000_0000_0000);
1481    }
1482}
leb128_tokio/lib.rs

leb128_tokio/
lib.rs
