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
84impl StreamFrame {
85 #[must_use]
88 pub fn from_u8(v: u8) -> Option<Self> {
89 match v {
90 0 => Some(Self::Open),
91 1 => Some(Self::OpenAck),
92 2 => Some(Self::Data),
93 3 => Some(Self::Credit),
94 4 => Some(Self::Close),
95 5 => Some(Self::CloseAck),
96 6 => Some(Self::Reset),
97 _ => None,
98 }
99 }
100
101 #[must_use]
103 pub fn as_u8(self) -> u8 {
104 self as u8
105 }
106}
107
108#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
111pub struct SealedPayload {
112 pub kem_enc: Bytes48,
115 pub ciphertext: Vec<u8>,
117}
118
119#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
125pub struct InnerMessage {
126 pub message_type: u32,
128 pub correlation_id: Bytes32,
130 pub compression: u8,
132 pub uncompressed_len: u32,
134 pub counter: u64,
136 pub timestamp_ms: u64,
138 pub expires_at: u64,
140 pub payload: Vec<u8>,
142 pub sender_sig: Bytes96,
144}
145
146#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
149pub struct DigMessageEnvelope {
150 pub version: u8,
152 pub message_type: u32,
154 pub flags: u8,
156 pub correlation_id: Bytes32,
158 pub sender: Bytes32,
160 pub recipient: Bytes32,
162 pub sender_epoch: u32,
164 pub stream: Option<StreamHeader>,
166 pub sealed: SealedPayload,
168}
169
170impl DigMessageEnvelope {
171 pub fn header_bytes(&self) -> Result<Vec<u8>> {
178 let mut out = Vec::new();
179 let codec = |e: chia_traits::Error| MessageError::Codec(e.to_string());
180 self.version.stream(&mut out).map_err(codec)?;
181 self.message_type.stream(&mut out).map_err(codec)?;
182 self.flags.stream(&mut out).map_err(codec)?;
183 self.correlation_id.stream(&mut out).map_err(codec)?;
184 self.sender.stream(&mut out).map_err(codec)?;
185 self.recipient.stream(&mut out).map_err(codec)?;
186 self.sender_epoch.stream(&mut out).map_err(codec)?;
187 self.stream.stream(&mut out).map_err(codec)?;
188 Ok(out)
189 }
190}
191
192pub fn encode_envelope(envelope: &DigMessageEnvelope) -> Result<Vec<u8>> {
198 let bytes = envelope
199 .to_bytes()
200 .map_err(|e| MessageError::Codec(e.to_string()))?;
201 if bytes.len() > MAX_ENVELOPE_BYTES {
202 return Err(MessageError::EnvelopeTooLarge {
203 size: bytes.len(),
204 max: MAX_ENVELOPE_BYTES,
205 });
206 }
207 Ok(bytes)
208}
209
210pub fn decode_envelope(bytes: &[u8]) -> Result<DigMessageEnvelope> {
217 if bytes.len() > MAX_ENVELOPE_BYTES {
218 return Err(MessageError::EnvelopeTooLarge {
219 size: bytes.len(),
220 max: MAX_ENVELOPE_BYTES,
221 });
222 }
223 let envelope = DigMessageEnvelope::from_bytes(bytes)
224 .map_err(|e| MessageError::Truncated(e.to_string()))?;
225 if envelope.version > ENVELOPE_VERSION {
226 return Err(MessageError::UnsupportedVersion(envelope.version));
227 }
228 Ok(envelope)
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234 use crate::compression::COMPRESSION_NONE;
235
236 fn sample(shape: InteractionShape, stream: Option<StreamHeader>) -> DigMessageEnvelope {
238 DigMessageEnvelope {
239 version: ENVELOPE_VERSION,
240 message_type: 0x0000_0201,
241 flags: shape.as_bits() | FLAG_SEALED,
242 correlation_id: Bytes32::new([1u8; 32]),
243 sender: Bytes32::new([2u8; 32]),
244 recipient: Bytes32::new([3u8; 32]),
245 sender_epoch: 7,
246 stream,
247 sealed: SealedPayload {
248 kem_enc: Bytes48::new([4u8; 48]),
249 ciphertext: vec![9, 8, 7, 6, 5],
250 },
251 }
252 }
253
254 #[test]
255 fn envelope_round_trips_for_every_shape() {
256 let cases = [
257 (InteractionShape::OneShot, None),
258 (InteractionShape::Request, None),
259 (InteractionShape::Response, None),
260 (
261 InteractionShape::StreamFrame,
262 Some(StreamHeader {
263 frame: StreamFrame::Data as u8,
264 seq: 42,
265 window: 8,
266 }),
267 ),
268 ];
269 for (shape, stream) in cases {
270 let env = sample(shape, stream);
271 let bytes = encode_envelope(&env).unwrap();
272 let decoded = decode_envelope(&bytes).unwrap();
273 assert_eq!(env, decoded, "{shape:?} must round-trip");
274 }
275 }
276
277 #[test]
278 fn encoding_is_byte_deterministic() {
279 let env = sample(InteractionShape::OneShot, None);
280 assert_eq!(
281 encode_envelope(&env).unwrap(),
282 encode_envelope(&env).unwrap()
283 );
284 }
285
286 #[test]
287 fn inner_message_round_trips() {
288 let inner = InnerMessage {
289 message_type: 0x0000_0201,
290 correlation_id: Bytes32::new([1u8; 32]),
291 compression: COMPRESSION_NONE,
292 uncompressed_len: 5,
293 counter: 11,
294 timestamp_ms: 1_700_000_000_000,
295 expires_at: 0,
296 payload: vec![1, 2, 3, 4, 5],
297 sender_sig: Bytes96::new([0u8; 96]),
298 };
299 let bytes = inner.to_bytes().unwrap();
300 assert_eq!(InnerMessage::from_bytes(&bytes).unwrap(), inner);
301 }
302
303 #[test]
304 fn unknown_newer_version_is_rejected() {
305 let mut env = sample(InteractionShape::OneShot, None);
306 env.version = ENVELOPE_VERSION + 1;
307 let bytes = env.to_bytes().unwrap();
308 assert_eq!(
309 decode_envelope(&bytes).unwrap_err(),
310 MessageError::UnsupportedVersion(ENVELOPE_VERSION + 1)
311 );
312 }
313
314 #[test]
315 fn oversized_frame_is_rejected_before_decoding() {
316 let bytes = vec![0u8; MAX_ENVELOPE_BYTES + 1];
317 assert_eq!(
318 decode_envelope(&bytes).unwrap_err(),
319 MessageError::EnvelopeTooLarge {
320 size: MAX_ENVELOPE_BYTES + 1,
321 max: MAX_ENVELOPE_BYTES
322 }
323 );
324 }
325
326 #[test]
327 fn oversized_envelope_encode_is_rejected() {
328 let mut env = sample(InteractionShape::OneShot, None);
329 env.sealed.ciphertext = vec![0u8; MAX_ENVELOPE_BYTES + 1];
330 assert!(matches!(
331 encode_envelope(&env).unwrap_err(),
332 MessageError::EnvelopeTooLarge { .. }
333 ));
334 }
335
336 #[test]
337 fn truncated_frame_is_rejected() {
338 let env = sample(InteractionShape::OneShot, None);
339 let bytes = encode_envelope(&env).unwrap();
340 let err = decode_envelope(&bytes[..bytes.len() / 2]).unwrap_err();
341 assert!(matches!(err, MessageError::Truncated(_)));
342 }
343
344 #[test]
345 fn header_bytes_excludes_the_sealed_region() {
346 let env = sample(InteractionShape::OneShot, None);
347 let header = env.header_bytes().unwrap();
348 let full = env.to_bytes().unwrap();
349 assert!(full.starts_with(&header));
351 assert!(header.len() < full.len());
352 }
353
354 #[test]
355 fn flags_shape_helpers_round_trip() {
356 for shape in [
357 InteractionShape::OneShot,
358 InteractionShape::Request,
359 InteractionShape::Response,
360 InteractionShape::StreamFrame,
361 ] {
362 let flags = shape.as_bits() | FLAG_SEALED;
363 assert_eq!(InteractionShape::from_flags(flags), Some(shape));
364 assert_eq!(flags & FLAG_SEALED, FLAG_SEALED);
365 }
366 }
367
368 #[test]
369 fn message_type_round_trips() {
370 let mt = MessageType(0x1000_0000);
371 assert_eq!(
372 MessageType::from_bytes(&mt.to_bytes().unwrap()).unwrap(),
373 mt
374 );
375 }
376}