1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

use abomonation::{decode, encode, measure, Abomonation};
use bytes::BytesMut;
use serde::{Deserialize, Serialize};
use std::{
    fmt::Debug,
    io::{Error, ErrorKind},
};

use crate::communication::CommunicationError;

/// Wrapper around a deserialized message. The wrapper can either own the deserialized
/// message or store a reference to it.
pub enum DeserializedMessage<'a, T> {
    Ref(&'a T),
    Owned(T),
}

/// Trait automatically derived for all messages that derive `Serialize` and `Deserialize`.
pub trait Serializable<'a>: Sized {
    fn encode(&self) -> Result<BytesMut, CommunicationError>;
    fn decode(buf: &'a mut BytesMut) -> Result<DeserializedMessage<'a, Self>, CommunicationError>;
}

impl<'a, D> Serializable<'a> for D
where
    D: Debug + Clone + Send + Serialize + Deserialize<'a>,
{
    default fn encode(&self) -> Result<BytesMut, CommunicationError> {
        let serialized_msg = bincode::serialize(self).map_err(|e| CommunicationError::from(e))?;
        let serialized_msg: BytesMut = BytesMut::from(serialized_msg);
        Ok(serialized_msg)
    }

    default fn decode(
        buf: &'a mut BytesMut,
    ) -> Result<DeserializedMessage<'a, D>, CommunicationError> {
        let msg: D = bincode::deserialize(buf).map_err(|e| CommunicationError::from(e))?;
        Ok(DeserializedMessage::Owned(msg))
    }
}

/// Specialized version used when messages derive `Abomonation`.
impl<'a, D> Serializable<'a> for D
where
    D: Debug + Clone + Send + Serialize + Deserialize<'a> + Abomonation,
{
    fn encode(&self) -> Result<BytesMut, CommunicationError> {
        let mut serialized_msg: Vec<u8> = Vec::with_capacity(measure(self));
        unsafe {
            encode(self, &mut serialized_msg)
                .map_err(|e| CommunicationError::AbomonationError(e))?;
        }
        let serialized_msg: BytesMut = BytesMut::from(serialized_msg);
        Ok(serialized_msg)
    }

    fn decode(buf: &'a mut BytesMut) -> Result<DeserializedMessage<'a, D>, CommunicationError> {
        let (msg, _) = {
            unsafe {
                match decode::<D>(buf.as_mut()) {
                    Some(msg) => msg,
                    None => {
                        return Err(CommunicationError::AbomonationError(Error::new(
                            ErrorKind::Other,
                            "Deserialization failed",
                        )))
                    }
                }
            }
        };
        Ok(DeserializedMessage::Ref(msg))
    }
}