tari_comms 6.2.0-pre.0

// Copyright 2019, The Tari Project
//
// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
// following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
// disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
// following disclaimer in the documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
// products derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::{
    fmt,
    fmt::{Display, Formatter},
};

use super::MessageError;
use crate::decode_budget::{DecodeBudget, decode_with_max_items};
// Re-export protos
pub use crate::proto::envelope::*;

/// Wraps a number of `prost::Message`s in a EnvelopeBody
#[macro_export]
macro_rules! wrap_in_envelope_body {
    ($($e:expr),+) => {{
        use $crate::message::MessageExt;
        let mut envelope_body = $crate::message::EnvelopeBody::new();
        $(
            envelope_body.push_part($e.to_encoded_bytes());
        )*
        envelope_body
    }}
}

impl EnvelopeBody {
    /// New empty envelope body.
    pub fn new() -> Self {
        Self {
            parts: Default::default(),
        }
    }

    /// Number of parts contained within this envelope.
    pub fn len(&self) -> usize {
        self.parts.len()
    }

    /// Total size of all parts contained within this envelope.
    pub fn total_size(&self) -> usize {
        self.parts.iter().fold(0, |acc: usize, b| acc.saturating_add(b.len()))
    }

    /// Returns true if the envelope is empty, otherwise false.
    pub fn is_empty(&self) -> bool {
        self.parts.is_empty()
    }

    /// Removes and returns the part at the given index. None
    /// is returned if the index is out of bounds
    pub fn take_part(&mut self, index: usize) -> Option<Vec<u8>> {
        Some(index)
            .filter(|i| self.parts.len() > *i)
            .map(|i| self.parts.remove(i))
    }

    /// Push a new part to the end of the envelope.
    pub fn push_part(&mut self, part: Vec<u8>) {
        self.parts.push(part)
    }

    /// Returns a Vec of message blobs.
    pub fn into_inner(self) -> Vec<Vec<u8>> {
        self.parts
    }

    /// Decodes a part of the message body and returns the result. If the part index is out of range Ok(None) is
    /// returned
    #[deprecated(note = "use the _with_max_items variant, which bounds what decoding can allocate")]
    pub fn decode_part<T>(&self, index: usize) -> Result<Option<T>, MessageError>
    where T: prost::Message + Default {
        match self.parts.get(index) {
            Some(part) => T::decode(part.as_slice()).map(Some).map_err(Into::into),
            None => Ok(None),
        }
    }

    /// Like [Self::decode_part], but rejects a part carrying more than `max_items` embedded items (see
    /// [crate::decode_budget]) before decoding it. The rejection is a [MessageError::DecodeError].
    pub fn decode_part_with_max_items<T>(&self, index: usize, max_items: usize) -> Result<Option<T>, MessageError>
    where T: prost::Message + Default + DecodeBudget {
        match self.parts.get(index) {
            Some(part) => decode_with_max_items(part, max_items).map(Some).map_err(Into::into),
            None => Ok(None),
        }
    }
}

impl Display for EnvelopeBody {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{} byte(s), {} part(s)", self.total_size(), self.len())
    }
}

#[cfg(test)]
mod test {
    use prost::Message;

    use super::*;

    fn body_of_empty_parts(count: usize) -> EnvelopeBody {
        EnvelopeBody {
            parts: vec![Vec::new(); count],
        }
    }

    #[test]
    fn decode_part_with_max_items_rejects_an_over_budget_part_before_decoding_it() {
        let flood = body_of_empty_parts(2_000).encode_to_vec();
        let at_budget = body_of_empty_parts(1_000).encode_to_vec();
        let body = EnvelopeBody {
            parts: vec![flood, at_budget],
        };

        // Prost alone accepts the flood
        assert_eq!(
            EnvelopeBody::decode(body.parts.first().unwrap().as_slice())
                .unwrap()
                .len(),
            2_000
        );
        let err = body.decode_part_with_max_items::<EnvelopeBody>(0, 1_000).unwrap_err();
        assert!(err.to_string().contains("decode budget"), "unexpected error: {err}");

        let decoded = body
            .decode_part_with_max_items::<EnvelopeBody>(1, 1_000)
            .unwrap()
            .unwrap();
        assert_eq!(decoded.len(), 1_000);
        assert!(
            body.decode_part_with_max_items::<EnvelopeBody>(2, 1_000)
                .unwrap()
                .is_none()
        );
    }
}