bufmut_netext/
lib.rs

1use bytes::{Buf, BufMut};
2use std::net::{Ipv4Addr, Ipv6Addr};
3use std::{convert::TryInto, fmt};
4
5pub trait Codec: Sized {
6    fn decode<B: Buf>(buf: &mut B) -> Self;
7    fn encode<B: BufMut>(&self, buf: &mut B);
8}
9
10pub trait BufExt {
11    fn get<T: Codec>(&mut self) -> T;
12    fn get_variable_length(&mut self) -> u64;
13}
14
15impl Codec for u8 {
16    fn decode<B: Buf>(buf: &mut B) -> u8 {
17        buf.get_u8()
18    }
19    fn encode<B: BufMut>(&self, buf: &mut B) {
20        buf.put_u8(*self);
21    }
22}
23
24impl Codec for u16 {
25    fn decode<B: Buf>(buf: &mut B) -> u16 {
26        buf.get_u16()
27    }
28    fn encode<B: BufMut>(&self, buf: &mut B) {
29        buf.put_u16(*self);
30    }
31}
32
33impl Codec for u32 {
34    fn decode<B: Buf>(buf: &mut B) -> u32 {
35        buf.get_u32()
36    }
37    fn encode<B: BufMut>(&self, buf: &mut B) {
38        buf.put_u32(*self);
39    }
40}
41
42impl Codec for u64 {
43    fn decode<B: Buf>(buf: &mut B) -> u64 {
44        buf.get_u64()
45    }
46    fn encode<B: BufMut>(&self, buf: &mut B) {
47        buf.put_u64(*self);
48    }
49}
50
51impl Codec for Ipv4Addr {
52    fn decode<B: Buf>(buf: &mut B) -> Ipv4Addr {
53        let mut octets = [0; 4];
54        buf.copy_to_slice(&mut octets);
55        octets.into()
56    }
57    fn encode<B: BufMut>(&self, buf: &mut B) {
58        buf.put_slice(&self.octets());
59    }
60}
61
62impl Codec for Ipv6Addr {
63    fn decode<B: Buf>(buf: &mut B) -> Ipv6Addr {
64        let mut octets = [0; 16];
65        buf.copy_to_slice(&mut octets);
66        octets.into()
67    }
68    fn encode<B: BufMut>(&self, buf: &mut B) {
69        buf.put_slice(&self.octets());
70    }
71}
72
73#[derive(Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
74pub struct VarInt(u64);
75
76impl VarInt {
77    pub const MAX: VarInt = VarInt((1 << 62) - 1);
78    pub const MAX_SIZE: usize = 0;
79
80    pub fn from_u32(x: u32) -> Self {
81        VarInt(x as u64)
82    }
83
84    pub fn from_u64(x: u64) -> Self {
85        VarInt(x)
86    }
87
88    pub fn into_inner(self) -> u64 {
89        self.0
90    }
91
92    pub fn size(self) -> usize {
93        let x = self.0;
94        if x < 2u64.pow(6) {
95            1
96        } else if x < 2u64.pow(14) {
97            2
98        } else if x < 2u64.pow(30) {
99            4
100        } else if x < 2u64.pow(62) {
101            8
102        } else {
103            unreachable!("VarInt malformed, size not computed");
104        }
105    }
106
107    pub fn size_encoded(first: u8) -> usize {
108        2usize.pow((first >> 6) as u32)
109    }
110}
111
112impl From<VarInt> for u64 {
113    fn from(x: VarInt) -> u64 {
114        x.0
115    }
116}
117
118impl From<u8> for VarInt {
119    fn from(x: u8) -> Self {
120        VarInt(x.into())
121    }
122}
123
124impl From<u16> for VarInt {
125    fn from(x: u16) -> Self {
126        VarInt(x.into())
127    }
128}
129
130impl From<u32> for VarInt {
131    fn from(x: u32) -> Self {
132        VarInt(x.into())
133    }
134}
135
136impl fmt::Debug for VarInt {
137    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138        self.0.fmt(f)
139    }
140}
141
142impl fmt::Display for VarInt {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        self.0.fmt(f)
145    }
146}
147
148impl Codec for VarInt {
149    fn decode<B: Buf>(r: &mut B) -> Self {
150        let mut buf = [0; 8];
151        buf[0] = r.get_u8();
152        let tag = buf[0] >> 6;
153        buf[0] &= 0b0011_1111;
154        let x = match tag {
155            0b00 => u64::from(buf[0]),
156            0b01 => {
157                r.copy_to_slice(&mut buf[1..2]);
158                u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
159            }
160            0b10 => {
161                r.copy_to_slice(&mut buf[1..4]);
162                u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
163            }
164            0b11 => {
165                r.copy_to_slice(&mut buf[1..8]);
166                u64::from_be_bytes(buf)
167            }
168            _ => unreachable!(),
169        };
170        VarInt(x)
171    }
172
173    fn encode<B: BufMut>(&self, w: &mut B) {
174        let x = self.0;
175        if x < 2u64.pow(6) {
176            w.put_u8(x as u8);
177        } else if x < 2u64.pow(14) {
178            w.put_u16(0b01 << 14 | x as u16);
179        } else if x < 2u64.pow(30) {
180            w.put_u32(0b10 << 30 | x as u32);
181        } else if x < 2u64.pow(62) {
182            w.put_u64(0b11 << 62 | x);
183        } else {
184            unreachable!("VarInt malformed")
185        }
186    }
187}
bufmut_netext/lib.rs

bufmut_netext/
lib.rs
