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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use core::ops::{Deref, DerefMut};
use ockam_core::{
    compat::{collections::BTreeMap, string::String, vec::Vec},
    Decodable, Encodable, Message, Result,
};
use serde::{Deserialize, Serialize};

/// A message metadata wrapper type
///
/// This message wraps around a well-typed Message type, with
/// additional metadata.  Metadata is split between the "scope" and
/// "generic" sections.
///
/// ## Scope metadata
///
/// This metadata is passed around in a particular metadata scope.
/// For example, a worker that adds some behaviour to message sending
/// may chose to embed "scope" metadata.  When wrapping this message
/// in another scope the previously scoped metadata becomes part of
/// the opaque `data` section.
///
/// Thus it is not possible to retrieve metadata from a different
/// nested scope!
///
/// ## Generic metadata
///
/// When creating an `OckamMessage` it's also possible to attach
/// generic metadata.  This data is passed around for every nested
/// scope and must be re-attached to the outest-most scope when
/// peeling a nested message stack.
#[derive(Message, Serialize, Deserialize)]
#[non_exhaustive]
pub struct OckamMessage {
    /// Main data section of this message
    pub data: Vec<u8>,
    /// Metadata for this specific scope
    pub scope: Vec<Vec<u8>>,
    /// Metadata that is carried to the final recipient of the message
    pub generic: Option<Metadata>,
}

impl OckamMessage {
    pub fn new<M: Message>(msg: M) -> Result<Self> {
        Ok(Self {
            data: msg.encode()?,
            scope: vec![],
            generic: None,
        })
    }

    /// Create a new `OckamMessage` by nesting a previous one
    pub fn wrap(mut prev: Self) -> Result<Self> {
        let generic = core::mem::replace(&mut prev.generic, None);
        Ok(Self {
            data: prev.encode()?,
            scope: vec![],
            generic,
        })
    }

    /// Add some metadata to this scope
    pub fn scope_data(mut self, meta: Vec<u8>) -> Self {
        self.scope.push(meta);
        self
    }

    /// Add to the generic metadata section
    pub fn generic_data<S: Into<String>>(mut self, key: S, val: Vec<u8>) -> Self {
        if self.generic.is_none() {
            self.generic = Some(Metadata(BTreeMap::new()));
        }

        self.generic.as_mut().unwrap().insert(key.into(), val);
        self
    }

    /// Dissolve this outer layer of Message and reveal nested message
    ///
    /// Will throw a deserialisation error if the inner data is NOT an
    /// OckamMessage!
    pub fn peel(mut self) -> Result<Self> {
        let generic = core::mem::replace(&mut self.generic, None);
        let mut peeled = Self::decode(&self.data)?;
        peeled.generic = generic;
        Ok(peeled)
    }
}

/// An encoding for message metadata
#[derive(Debug, Serialize, Deserialize)]
pub struct Metadata(BTreeMap<String, Vec<u8>>);

impl Deref for Metadata {
    type Target = BTreeMap<String, Vec<u8>>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Metadata {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// This test emulates the message flow for a very simple pipe
/// metadata message.  At the core we have some piece of data, which
/// gets wrapped in a bespoke message type.  This message then is
/// attached with scope metadata to indicate the message index.
#[test]
fn nest_metadata() {
    #[derive(Serialize, Deserialize, Message, PartialEq, Debug, Clone)]
    struct FakePipeMsg {
        vec: Vec<u8>,
    }

    let base_msg = FakePipeMsg {
        vec: vec![1, 2, 3, 4, 5, 6, 7, 8],
    };

    // The base message type with a msg_type generic metadata field
    let ockam_msg1 = OckamMessage::new(base_msg.clone())
        .unwrap()
        .generic_data("msg_type", "pipemsg".as_bytes().into());

    // Wrap this message in another scope which adds a message index
    // to the scope metadata section.  `vec![1]` is our index here but
    // in reality this should be a properly encoded type!
    let ockam_msg2 = OckamMessage::wrap(ockam_msg1).unwrap().scope_data(vec![1]);

    /////////// On the other side of a transport

    let msg_type = ockam_msg2
        .generic
        .as_ref()
        .unwrap()
        .get("msg_type")
        .unwrap();
    assert_eq!(msg_type, "pipemsg".as_bytes());

    let index = ockam_msg2.scope.get(0).unwrap();
    assert_eq!(index, &vec![1]);

    // Then we peel the previous message type
    let ockam_msg1 = ockam_msg2.peel().unwrap();
    let base_msg_other_side = FakePipeMsg::decode(&ockam_msg1.data).unwrap();

    assert_eq!(base_msg, base_msg_other_side);
}