1use std::net::{Ipv4Addr, Ipv6Addr};
11
12use bytes::{Buf, BufMut};
13use thiserror::Error;
14
15use crate::VarInt;
16
17#[derive(Error, Debug, Copy, Clone, Eq, PartialEq)]
19#[error("unexpected end of buffer")]
20pub struct UnexpectedEnd;
21
22pub type Result<T> = ::std::result::Result<T, UnexpectedEnd>;
24
25pub trait Codec: Sized {
27 fn decode<B: Buf>(buf: &mut B) -> Result<Self>;
29 fn encode<B: BufMut>(&self, buf: &mut B);
31}
32
33impl Codec for u8 {
34 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
35 if buf.remaining() < 1 {
36 return Err(UnexpectedEnd);
37 }
38 Ok(buf.get_u8())
39 }
40 fn encode<B: BufMut>(&self, buf: &mut B) {
41 buf.put_u8(*self);
42 }
43}
44
45impl Codec for u16 {
46 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
47 if buf.remaining() < 2 {
48 return Err(UnexpectedEnd);
49 }
50 Ok(buf.get_u16())
51 }
52 fn encode<B: BufMut>(&self, buf: &mut B) {
53 buf.put_u16(*self);
54 }
55}
56
57impl Codec for u32 {
58 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
59 if buf.remaining() < 4 {
60 return Err(UnexpectedEnd);
61 }
62 Ok(buf.get_u32())
63 }
64 fn encode<B: BufMut>(&self, buf: &mut B) {
65 buf.put_u32(*self);
66 }
67}
68
69impl Codec for u64 {
70 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
71 if buf.remaining() < 8 {
72 return Err(UnexpectedEnd);
73 }
74 Ok(buf.get_u64())
75 }
76 fn encode<B: BufMut>(&self, buf: &mut B) {
77 buf.put_u64(*self);
78 }
79}
80
81impl Codec for Ipv4Addr {
82 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
83 if buf.remaining() < 4 {
84 return Err(UnexpectedEnd);
85 }
86 let mut octets = [0; 4];
87 buf.copy_to_slice(&mut octets);
88 Ok(octets.into())
89 }
90 fn encode<B: BufMut>(&self, buf: &mut B) {
91 buf.put_slice(&self.octets());
92 }
93}
94
95impl Codec for Ipv6Addr {
96 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
97 if buf.remaining() < 16 {
98 return Err(UnexpectedEnd);
99 }
100 let mut octets = [0; 16];
101 buf.copy_to_slice(&mut octets);
102 Ok(octets.into())
103 }
104 fn encode<B: BufMut>(&self, buf: &mut B) {
105 buf.put_slice(&self.octets());
106 }
107}
108
109pub trait BufExt {
111 fn get<T: Codec>(&mut self) -> Result<T>;
113 fn get_var(&mut self) -> Result<u64>;
115}
116
117impl<T: Buf> BufExt for T {
118 fn get<U: Codec>(&mut self) -> Result<U> {
119 U::decode(self)
120 }
121
122 fn get_var(&mut self) -> Result<u64> {
123 Ok(VarInt::decode(self)?.into_inner())
124 }
125}
126
127pub trait BufMutExt {
129 fn write<T: Codec>(&mut self, x: T);
131 fn write_var(&mut self, x: u64);
133}
134
135impl<T: BufMut> BufMutExt for T {
136 fn write<U: Codec>(&mut self, x: U) {
137 x.encode(self);
138 }
139
140 fn write_var(&mut self, x: u64) {
141 VarInt::from_u64(x).unwrap().encode(self);
142 }
143}