sheave_core/messages/
peer_bandwidth.rs

1mod limit_type;
2
3use std::{
4    cmp::Ordering,
5    io::Result as IOResult,
6    ops::Div
7};
8use crate::{
9    ByteBuffer,
10    Decoder,
11    Encoder,
12    messages::{
13        Channel,
14        ChunkData,
15        headers::MessageType
16    }
17};
18pub use self::limit_type::*;
19
20/// The message to tell the client-side bandwidth.
21/// This has 2 ways of comparision by which field you specifies.
22///
23/// # Examples
24///
25/// ```rust
26/// use sheave_core::messages::{
27///     LimitType,
28///     PeerBandwidth,
29/// };
30///
31/// let peer_bandwidth = PeerBandwidth::default();
32///
33/// // When you compare this message with a bandwidth number.
34/// assert!(2500000u32 == peer_bandwidth);
35/// assert!(0 <= peer_bandwidth);
36/// // When you compare this message with a limit type.
37/// assert!(LimitType::default() == peer_bandwidth)
38/// ```
39#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
40pub struct PeerBandwidth(u32, LimitType);
41
42impl PeerBandwidth {
43    const DEFAULT: u32 = 2500000;
44
45    /// Constructs a PeerBandwidth message.
46    pub fn new(peer_bandwidth: u32, limit_type: LimitType) -> Self {
47        Self(peer_bandwidth, limit_type)
48    }
49
50    pub fn get_inner_bandwidth(&self) -> u32 {
51        self.0
52    }
53
54    pub fn get_inner_limit_type(&self) -> LimitType {
55        self.1
56    }
57}
58
59impl Default for PeerBandwidth {
60    /// Constructs a PeerBandwidth message with the default bandwidth and the default limit type (2500000 in bits, 2 (Dynamic)).
61    ///
62    /// # Examples
63    ///
64    /// ```rust
65    /// use sheave_core::messages::{
66    ///     LimitType,
67    ///     PeerBandwidth
68    /// };
69    ///
70    /// let peer_bandwidth = PeerBandwidth::default();
71    /// assert_eq!(2500000u32, peer_bandwidth);
72    /// assert_eq!(LimitType::default(), peer_bandwidth)
73    /// ```
74    fn default() -> Self {
75        Self(Self::DEFAULT, LimitType::default())
76    }
77}
78
79impl PartialEq<u32> for PeerBandwidth {
80    fn eq(&self, other: &u32) -> bool {
81        self.0.eq(other)
82    }
83}
84
85impl PartialEq<PeerBandwidth> for u32 {
86    fn eq(&self, other: &PeerBandwidth) -> bool {
87        self.eq(&other.0)
88    }
89}
90
91impl PartialOrd<u32> for PeerBandwidth {
92    fn partial_cmp(&self, other: &u32) -> Option<Ordering> {
93        self.0.partial_cmp(other)
94    }
95}
96
97impl PartialOrd<PeerBandwidth> for u32 {
98    fn partial_cmp(&self, other: &PeerBandwidth) -> Option<Ordering> {
99        self.partial_cmp(&other.0)
100    }
101}
102
103impl PartialEq<LimitType> for PeerBandwidth {
104    fn eq(&self, other: &LimitType) -> bool {
105        self.1.eq(other)
106    }
107}
108
109impl PartialEq<PeerBandwidth> for LimitType {
110    fn eq(&self, other: &PeerBandwidth) -> bool {
111        self.eq(&other.1)
112    }
113}
114
115impl Div<u32> for PeerBandwidth {
116    type Output = Self;
117
118    fn div(self, rhs: u32) -> Self::Output {
119        Self(self.0 / rhs, self.1)
120    }
121}
122
123impl ChunkData for PeerBandwidth {
124    const CHANNEL: Channel = Channel::Network;
125    const MESSAGE_TYPE: MessageType = MessageType::PeerBandwidth;
126}
127
128impl Decoder<PeerBandwidth> for ByteBuffer {
129    /// Decodes bytes into a PeerBandwidth message.
130    ///
131    /// # Errors
132    ///
133    /// * [`InsufficientBufferLength`]
134    ///
135    /// When this buffer didn't remain at least 5 bytes.
136    ///
137    /// # Examples
138    ///
139    /// ```rust
140    /// use sheave_core::{
141    ///     ByteBuffer,
142    ///     Decoder,
143    ///     messages::PeerBandwidth
144    /// };
145    ///
146    /// let mut buffer = ByteBuffer::default();
147    /// buffer.put_u32_be(u32::default());
148    /// buffer.put_u8(u8::default());
149    /// assert!(Decoder::<PeerBandwidth>::decode(&mut buffer).is_ok())
150    /// ```
151    ///
152    /// [`InsufficientBufferLength`]: crate::byte_buffer::InsufficientBufferLength
153    fn decode(&mut self) -> IOResult<PeerBandwidth> {
154        let bandwidth = self.get_u32_be()?;
155        let limit_type = self.get_u8()?;
156        Ok(PeerBandwidth(bandwidth, limit_type.into()))
157    }
158}
159
160impl Encoder<PeerBandwidth> for ByteBuffer {
161    /// Encodes a PeerBandwidth message into bytes.
162    fn encode(&mut self, peer_bandwidth: &PeerBandwidth) {
163        self.put_u32_be(peer_bandwidth.0);
164        self.put_u8(peer_bandwidth.1.into());
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[test]
173    fn decode_peer_bandwidth() {
174        let mut buffer = ByteBuffer::default();
175        buffer.put_u32_be(2500000);
176        buffer.put_u8(2);
177        let result: IOResult<PeerBandwidth> = buffer.decode();
178        assert!(result.is_ok());
179        let actual = result.unwrap();
180        let expected = PeerBandwidth::default();
181        assert_eq!(expected, actual)
182    }
183
184    #[test]
185    fn encode_peer_bandwidth() {
186        let mut buffer = ByteBuffer::default();
187        let expected_bandwidth = 2500000u32;
188        let expected_limit_type = 2u8;
189        buffer.encode(&PeerBandwidth::default());
190        let actual_bandwidth = buffer.get_u32_be().unwrap();
191        assert_eq!(expected_bandwidth, actual_bandwidth);
192        let actual_limit_type = buffer.get_u8().unwrap();
193        assert_eq!(expected_limit_type, actual_limit_type)
194    }
195}