1use chia_protocol::{Bytes32, Bytes48, Bytes96};
10use chia_streamable_macro::Streamable;
11use chia_traits::Streamable as StreamableTrait;
12
13use crate::constants::{ENVELOPE_VERSION, MAX_ENVELOPE_BYTES};
14use crate::error::{MessageError, Result};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
19pub struct MessageType(pub u32);
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum InteractionShape {
24 OneShot = 0,
26 Request = 1,
28 Response = 2,
30 StreamFrame = 3,
32}
33
34pub const FLAG_SHAPE_MASK: u8 = 0b0000_0011;
36pub const FLAG_SEALED: u8 = 0b0000_0100;
38
39impl InteractionShape {
40 #[must_use]
42 pub fn from_flags(flags: u8) -> Option<Self> {
43 match flags & FLAG_SHAPE_MASK {
44 0 => Some(Self::OneShot),
45 1 => Some(Self::Request),
46 2 => Some(Self::Response),
47 3 => Some(Self::StreamFrame),
48 _ => None,
49 }
50 }
51
52 #[must_use]
54 pub fn as_bits(self) -> u8 {
55 self as u8
56 }
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Streamable)]
62pub struct StreamHeader {
63 pub frame: u8,
65 pub seq: u64,
67 pub window: u32,
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub enum StreamFrame {
75 Open = 0,
76 OpenAck = 1,
77 Data = 2,
78 Credit = 3,
79 Close = 4,
80 CloseAck = 5,
81 Reset = 6,
82}
83
84#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
87pub struct SealedPayload {
88 pub kem_enc: Bytes48,
91 pub ciphertext: Vec<u8>,
93}
94
95#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
101pub struct InnerMessage {
102 pub message_type: u32,
104 pub correlation_id: Bytes32,
106 pub compression: u8,
108 pub uncompressed_len: u32,
110 pub counter: u64,
112 pub timestamp_ms: u64,
114 pub expires_at: u64,
116 pub payload: Vec<u8>,
118 pub sender_sig: Bytes96,
120}
121
122#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
125pub struct DigMessageEnvelope {
126 pub version: u8,
128 pub message_type: u32,
130 pub flags: u8,
132 pub correlation_id: Bytes32,
134 pub sender: Bytes32,
136 pub recipient: Bytes32,
138 pub sender_epoch: u32,
140 pub stream: Option<StreamHeader>,
142 pub sealed: SealedPayload,
144}
145
146impl DigMessageEnvelope {
147 pub fn header_bytes(&self) -> Result<Vec<u8>> {
154 let mut out = Vec::new();
155 let codec = |e: chia_traits::Error| MessageError::Codec(e.to_string());
156 self.version.stream(&mut out).map_err(codec)?;
157 self.message_type.stream(&mut out).map_err(codec)?;
158 self.flags.stream(&mut out).map_err(codec)?;
159 self.correlation_id.stream(&mut out).map_err(codec)?;
160 self.sender.stream(&mut out).map_err(codec)?;
161 self.recipient.stream(&mut out).map_err(codec)?;
162 self.sender_epoch.stream(&mut out).map_err(codec)?;
163 self.stream.stream(&mut out).map_err(codec)?;
164 Ok(out)
165 }
166}
167
168pub fn encode_envelope(envelope: &DigMessageEnvelope) -> Result<Vec<u8>> {
174 let bytes = envelope
175 .to_bytes()
176 .map_err(|e| MessageError::Codec(e.to_string()))?;
177 if bytes.len() > MAX_ENVELOPE_BYTES {
178 return Err(MessageError::EnvelopeTooLarge {
179 size: bytes.len(),
180 max: MAX_ENVELOPE_BYTES,
181 });
182 }
183 Ok(bytes)
184}
185
186pub fn decode_envelope(bytes: &[u8]) -> Result<DigMessageEnvelope> {
193 if bytes.len() > MAX_ENVELOPE_BYTES {
194 return Err(MessageError::EnvelopeTooLarge {
195 size: bytes.len(),
196 max: MAX_ENVELOPE_BYTES,
197 });
198 }
199 let envelope = DigMessageEnvelope::from_bytes(bytes)
200 .map_err(|e| MessageError::Truncated(e.to_string()))?;
201 if envelope.version > ENVELOPE_VERSION {
202 return Err(MessageError::UnsupportedVersion(envelope.version));
203 }
204 Ok(envelope)
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210 use crate::compression::COMPRESSION_NONE;
211
212 fn sample(shape: InteractionShape, stream: Option<StreamHeader>) -> DigMessageEnvelope {
214 DigMessageEnvelope {
215 version: ENVELOPE_VERSION,
216 message_type: 0x0000_0201,
217 flags: shape.as_bits() | FLAG_SEALED,
218 correlation_id: Bytes32::new([1u8; 32]),
219 sender: Bytes32::new([2u8; 32]),
220 recipient: Bytes32::new([3u8; 32]),
221 sender_epoch: 7,
222 stream,
223 sealed: SealedPayload {
224 kem_enc: Bytes48::new([4u8; 48]),
225 ciphertext: vec![9, 8, 7, 6, 5],
226 },
227 }
228 }
229
230 #[test]
231 fn envelope_round_trips_for_every_shape() {
232 let cases = [
233 (InteractionShape::OneShot, None),
234 (InteractionShape::Request, None),
235 (InteractionShape::Response, None),
236 (
237 InteractionShape::StreamFrame,
238 Some(StreamHeader {
239 frame: StreamFrame::Data as u8,
240 seq: 42,
241 window: 8,
242 }),
243 ),
244 ];
245 for (shape, stream) in cases {
246 let env = sample(shape, stream);
247 let bytes = encode_envelope(&env).unwrap();
248 let decoded = decode_envelope(&bytes).unwrap();
249 assert_eq!(env, decoded, "{shape:?} must round-trip");
250 }
251 }
252
253 #[test]
254 fn encoding_is_byte_deterministic() {
255 let env = sample(InteractionShape::OneShot, None);
256 assert_eq!(
257 encode_envelope(&env).unwrap(),
258 encode_envelope(&env).unwrap()
259 );
260 }
261
262 #[test]
263 fn inner_message_round_trips() {
264 let inner = InnerMessage {
265 message_type: 0x0000_0201,
266 correlation_id: Bytes32::new([1u8; 32]),
267 compression: COMPRESSION_NONE,
268 uncompressed_len: 5,
269 counter: 11,
270 timestamp_ms: 1_700_000_000_000,
271 expires_at: 0,
272 payload: vec![1, 2, 3, 4, 5],
273 sender_sig: Bytes96::new([0u8; 96]),
274 };
275 let bytes = inner.to_bytes().unwrap();
276 assert_eq!(InnerMessage::from_bytes(&bytes).unwrap(), inner);
277 }
278
279 #[test]
280 fn unknown_newer_version_is_rejected() {
281 let mut env = sample(InteractionShape::OneShot, None);
282 env.version = ENVELOPE_VERSION + 1;
283 let bytes = env.to_bytes().unwrap();
284 assert_eq!(
285 decode_envelope(&bytes).unwrap_err(),
286 MessageError::UnsupportedVersion(ENVELOPE_VERSION + 1)
287 );
288 }
289
290 #[test]
291 fn oversized_frame_is_rejected_before_decoding() {
292 let bytes = vec![0u8; MAX_ENVELOPE_BYTES + 1];
293 assert_eq!(
294 decode_envelope(&bytes).unwrap_err(),
295 MessageError::EnvelopeTooLarge {
296 size: MAX_ENVELOPE_BYTES + 1,
297 max: MAX_ENVELOPE_BYTES
298 }
299 );
300 }
301
302 #[test]
303 fn oversized_envelope_encode_is_rejected() {
304 let mut env = sample(InteractionShape::OneShot, None);
305 env.sealed.ciphertext = vec![0u8; MAX_ENVELOPE_BYTES + 1];
306 assert!(matches!(
307 encode_envelope(&env).unwrap_err(),
308 MessageError::EnvelopeTooLarge { .. }
309 ));
310 }
311
312 #[test]
313 fn truncated_frame_is_rejected() {
314 let env = sample(InteractionShape::OneShot, None);
315 let bytes = encode_envelope(&env).unwrap();
316 let err = decode_envelope(&bytes[..bytes.len() / 2]).unwrap_err();
317 assert!(matches!(err, MessageError::Truncated(_)));
318 }
319
320 #[test]
321 fn header_bytes_excludes_the_sealed_region() {
322 let env = sample(InteractionShape::OneShot, None);
323 let header = env.header_bytes().unwrap();
324 let full = env.to_bytes().unwrap();
325 assert!(full.starts_with(&header));
327 assert!(header.len() < full.len());
328 }
329
330 #[test]
331 fn flags_shape_helpers_round_trip() {
332 for shape in [
333 InteractionShape::OneShot,
334 InteractionShape::Request,
335 InteractionShape::Response,
336 InteractionShape::StreamFrame,
337 ] {
338 let flags = shape.as_bits() | FLAG_SEALED;
339 assert_eq!(InteractionShape::from_flags(flags), Some(shape));
340 assert_eq!(flags & FLAG_SEALED, FLAG_SEALED);
341 }
342 }
343
344 #[test]
345 fn message_type_round_trips() {
346 let mt = MessageType(0x1000_0000);
347 assert_eq!(
348 MessageType::from_bytes(&mt.to_bytes().unwrap()).unwrap(),
349 mt
350 );
351 }
352}