saorsa_transport/
coding.rs1use std::net::{Ipv4Addr, Ipv6Addr};
16
17use bytes::{Buf, BufMut};
18use thiserror::Error;
19
20use crate::{VarInt, VarIntBoundsExceeded};
21
22#[derive(Error, Debug, Copy, Clone, Eq, PartialEq)]
24#[error("unexpected end of buffer")]
25pub struct UnexpectedEnd;
26
27pub type Result<T> = ::std::result::Result<T, UnexpectedEnd>;
29
30pub trait Codec: Sized {
32 fn decode<B: Buf>(buf: &mut B) -> Result<Self>;
34 fn encode<B: BufMut>(&self, buf: &mut B);
36}
37
38impl Codec for u8 {
39 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
40 if buf.remaining() < 1 {
41 return Err(UnexpectedEnd);
42 }
43 Ok(buf.get_u8())
44 }
45 fn encode<B: BufMut>(&self, buf: &mut B) {
46 buf.put_u8(*self);
47 }
48}
49
50impl Codec for u16 {
51 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
52 if buf.remaining() < 2 {
53 return Err(UnexpectedEnd);
54 }
55 Ok(buf.get_u16())
56 }
57 fn encode<B: BufMut>(&self, buf: &mut B) {
58 buf.put_u16(*self);
59 }
60}
61
62impl Codec for u32 {
63 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
64 if buf.remaining() < 4 {
65 return Err(UnexpectedEnd);
66 }
67 Ok(buf.get_u32())
68 }
69 fn encode<B: BufMut>(&self, buf: &mut B) {
70 buf.put_u32(*self);
71 }
72}
73
74impl Codec for u64 {
75 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
76 if buf.remaining() < 8 {
77 return Err(UnexpectedEnd);
78 }
79 Ok(buf.get_u64())
80 }
81 fn encode<B: BufMut>(&self, buf: &mut B) {
82 buf.put_u64(*self);
83 }
84}
85
86impl Codec for Ipv4Addr {
87 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
88 if buf.remaining() < 4 {
89 return Err(UnexpectedEnd);
90 }
91 let mut octets = [0; 4];
92 buf.copy_to_slice(&mut octets);
93 Ok(octets.into())
94 }
95 fn encode<B: BufMut>(&self, buf: &mut B) {
96 buf.put_slice(&self.octets());
97 }
98}
99
100impl Codec for Ipv6Addr {
101 fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
102 if buf.remaining() < 16 {
103 return Err(UnexpectedEnd);
104 }
105 let mut octets = [0; 16];
106 buf.copy_to_slice(&mut octets);
107 Ok(octets.into())
108 }
109 fn encode<B: BufMut>(&self, buf: &mut B) {
110 buf.put_slice(&self.octets());
111 }
112}
113
114pub trait BufExt {
116 fn get<T: Codec>(&mut self) -> Result<T>;
118 fn get_var(&mut self) -> Result<u64>;
120}
121
122impl<T: Buf> BufExt for T {
123 fn get<U: Codec>(&mut self) -> Result<U> {
124 U::decode(self)
125 }
126
127 fn get_var(&mut self) -> Result<u64> {
128 Ok(VarInt::decode(self)?.into_inner())
129 }
130}
131
132pub trait BufMutExt {
134 fn write<T: Codec>(&mut self, x: T);
136 fn write_var(&mut self, x: u64) -> std::result::Result<(), VarIntBoundsExceeded>;
138 fn write_var_or_debug_assert(&mut self, x: u64) {
140 if self.write_var(x).is_err() {
141 tracing::error!("VarInt overflow: {} exceeds maximum", x);
142 debug_assert!(false, "VarInt overflow: {}", x);
143 }
144 }
145}
146
147impl<T: BufMut> BufMutExt for T {
148 fn write<U: Codec>(&mut self, x: U) {
149 x.encode(self);
150 }
151
152 fn write_var(&mut self, x: u64) -> std::result::Result<(), VarIntBoundsExceeded> {
153 VarInt::encode_checked(x, self)
154 }
155}