Skip to main content

saorsa_transport/
coding.rs

1// Copyright 2024 Saorsa Labs Ltd.
2//
3// This software is licensed under the MIT license <LICENSE-MIT or
4// https://opensource.org/licenses/MIT> or the Apache License, Version 2.0
5// <LICENSE-APACHE or https://www.apache.org/licenses/LICENSE-2.0>, at your
6// option. This file may not be copied, modified, or distributed except
7// according to those terms.
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under these licenses is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
13//! Coding related traits.
14
15use std::net::{Ipv4Addr, Ipv6Addr};
16
17use bytes::{Buf, BufMut};
18use thiserror::Error;
19
20use crate::{VarInt, VarIntBoundsExceeded};
21
22/// Error indicating that the provided buffer was too small
23#[derive(Error, Debug, Copy, Clone, Eq, PartialEq)]
24#[error("unexpected end of buffer")]
25pub struct UnexpectedEnd;
26
27/// Coding result type
28pub type Result<T> = ::std::result::Result<T, UnexpectedEnd>;
29
30/// Infallible encoding and decoding of QUIC primitives
31pub trait Codec: Sized {
32    /// Decode a `Self` from the provided buffer, if the buffer is large enough
33    fn decode<B: Buf>(buf: &mut B) -> Result<Self>;
34    /// Append the encoding of `self` to the provided buffer
35    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
114/// Extension trait for reading from buffers
115pub trait BufExt {
116    /// Read and decode a value from the buffer
117    fn get<T: Codec>(&mut self) -> Result<T>;
118    /// Read a variable-length integer from the buffer
119    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
132/// Extension trait for writing to buffers
133pub trait BufMutExt {
134    /// Write and encode a value to the buffer
135    fn write<T: Codec>(&mut self, x: T);
136    /// Write a variable-length integer to the buffer
137    fn write_var(&mut self, x: u64) -> std::result::Result<(), VarIntBoundsExceeded>;
138    /// Write a variable-length integer, debug-asserting on overflow in debug builds
139    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}