amaru_protocols/keepalive/
messages.rs

1// Copyright 2025 PRAGMA
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use amaru_kernel::cbor;
16
17#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
18pub struct Cookie(u16);
19
20impl Default for Cookie {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl Cookie {
27    pub fn new() -> Self {
28        Self(0)
29    }
30
31    pub fn next(self) -> Self {
32        Self(self.0.wrapping_add(1))
33    }
34
35    pub fn as_u16(self) -> u16 {
36        self.0
37    }
38}
39
40impl From<u16> for Cookie {
41    fn from(value: u16) -> Self {
42        Self(value)
43    }
44}
45
46impl From<Cookie> for u16 {
47    fn from(value: Cookie) -> Self {
48        value.0
49    }
50}
51
52impl<T> cbor::Encode<T> for Cookie {
53    fn encode<W: cbor::encode::Write>(
54        &self,
55        e: &mut cbor::Encoder<W>,
56        _ctx: &mut T,
57    ) -> Result<(), cbor::encode::Error<W::Error>> {
58        e.u16(self.0)?;
59        Ok(())
60    }
61}
62
63impl<'b, T> cbor::Decode<'b, T> for Cookie {
64    fn decode(d: &mut cbor::Decoder<'b>, _ctx: &mut T) -> Result<Self, cbor::decode::Error> {
65        Ok(Self(d.u16()?))
66    }
67}
68
69#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
70pub enum Message {
71    KeepAlive(Cookie),
72    ResponseKeepAlive(Cookie),
73    Done,
74}
75
76impl Message {
77    pub fn message_type(&self) -> &str {
78        match self {
79            Message::KeepAlive(_) => "KeepAlive",
80            Message::ResponseKeepAlive(_) => "ResponseKeepAlive",
81            Message::Done => "Done",
82        }
83    }
84}
85
86impl<T> cbor::Encode<T> for Message {
87    fn encode<W: cbor::encode::Write>(
88        &self,
89        e: &mut cbor::Encoder<W>,
90        _ctx: &mut T,
91    ) -> Result<(), cbor::encode::Error<W::Error>> {
92        match self {
93            Message::KeepAlive(cookie) => {
94                e.array(2)?.u16(0)?;
95                e.encode(cookie)?;
96            }
97            Message::ResponseKeepAlive(cookie) => {
98                e.array(2)?.u16(1)?;
99                e.encode(cookie)?;
100            }
101            Message::Done => {
102                e.array(1)?.u16(2)?;
103            }
104        }
105
106        Ok(())
107    }
108}
109
110impl<'b, T> cbor::Decode<'b, T> for Message {
111    fn decode(d: &mut cbor::Decoder<'b>, _ctx: &mut T) -> Result<Self, cbor::decode::Error> {
112        let len = d.array()?;
113        let label = d.u16()?;
114
115        match label {
116            0 => {
117                cbor::check_tagged_array_length(0, len, 2)?;
118                let cookie = d.decode()?;
119                Ok(Message::KeepAlive(cookie))
120            }
121            1 => {
122                cbor::check_tagged_array_length(1, len, 2)?;
123                let cookie = d.decode()?;
124                Ok(Message::ResponseKeepAlive(cookie))
125            }
126            2 => {
127                cbor::check_tagged_array_length(2, len, 1)?;
128                Ok(Message::Done)
129            }
130            _ => Err(cbor::decode::Error::message("can't decode Message")),
131        }
132    }
133}
134
135/// Roundtrip property tests for handshake messages.
136#[cfg(test)]
137mod tests {
138    use amaru_kernel::prop_cbor_roundtrip;
139    use proptest::{prelude::*, prop_compose};
140
141    use super::*;
142    use crate::keepalive::messages::Message::*;
143
144    prop_cbor_roundtrip!(Message, any_message());
145
146    // HELPERS
147
148    prop_compose! {
149        fn any_cookie()(n in any::<u16>()) -> Cookie {
150            Cookie(n)
151        }
152    }
153
154    prop_compose! {
155        fn any_keep_alive_message()(cookie in any_cookie()) -> Message {
156            KeepAlive(cookie)
157        }
158    }
159
160    prop_compose! {
161        fn any_response_keep_alive_message()(cookie in any_cookie()) -> Message {
162            ResponseKeepAlive(cookie)
163        }
164    }
165
166    pub fn done_message() -> impl Strategy<Value = Message> {
167        Just(Done)
168    }
169
170    pub fn any_message() -> impl Strategy<Value = Message> {
171        prop_oneof![
172            1 => done_message(),
173            1 => any_keep_alive_message(),
174            1 => any_response_keep_alive_message(),
175        ]
176    }
177}
amaru_protocols/keepalive/messages.rs

amaru_protocols/keepalive/
messages.rs
