ant_quic/
coding.rs

1// Copyright 2024 Saorsa Labs Ltd.
2//
3// This Saorsa Network Software is licensed under the General Public License (GPL), version 3.
4// Please see the file LICENSE-GPL, or visit <http://www.gnu.org/licenses/> for the full text.
5//
6// Full details available at https://saorsalabs.com/licenses
7
8//! Coding related traits.
9
10use std::net::{Ipv4Addr, Ipv6Addr};
11
12use bytes::{Buf, BufMut};
13use thiserror::Error;
14
15use crate::VarInt;
16
17/// Error indicating that the provided buffer was too small
18#[derive(Error, Debug, Copy, Clone, Eq, PartialEq)]
19#[error("unexpected end of buffer")]
20pub struct UnexpectedEnd;
21
22/// Coding result type
23pub type Result<T> = ::std::result::Result<T, UnexpectedEnd>;
24
25/// Infallible encoding and decoding of QUIC primitives
26pub trait Codec: Sized {
27    /// Decode a `Self` from the provided buffer, if the buffer is large enough
28    fn decode<B: Buf>(buf: &mut B) -> Result<Self>;
29    /// Append the encoding of `self` to the provided buffer
30    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
109/// Extension trait for reading from buffers
110pub trait BufExt {
111    /// Read and decode a value from the buffer
112    fn get<T: Codec>(&mut self) -> Result<T>;
113    /// Read a variable-length integer from the buffer
114    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
127/// Extension trait for writing to buffers
128pub trait BufMutExt {
129    /// Write and encode a value to the buffer
130    fn write<T: Codec>(&mut self, x: T);
131    /// Write a variable-length integer to the buffer
132    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}