Skip to main content

dig_message/
envelope.rs

1//! The base envelope + inner message — the byte-deterministic Chia-Streamable wire shapes (SPEC §2,
2//! §5.2) and the length-framed, size-bounded codec (SPEC §1).
3//!
4//! The envelope header (fields 1-8) is cleartext so a relay can route on `recipient` and multiplex on
5//! `correlation_id`/`stream`; ALL content lives in the sealed region (field 9). WU1 defines the shapes
6//! and the framing; the seal (`SealedPayload.kem_enc` + `ciphertext`) and the signature
7//! (`InnerMessage.sender_sig`) are populated by WU2 — their FIELDS are final here.
8
9use 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/// Extensible message-type id (SPEC §4). Additive-only: an id, once assigned, is never renumbered or
17/// repurposed. The runtime registry that dispatches on it is WU3; the wire newtype is defined here.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
19pub struct MessageType(pub u32);
20
21/// The interaction shape carried in `flags` bits 0-1 (SPEC §2 field 3 / §3).
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum InteractionShape {
24    /// Fire-and-forget: a single envelope, no response expected.
25    OneShot = 0,
26    /// A correlated request awaiting a response.
27    Request = 1,
28    /// A correlated response echoing a request's `correlation_id`.
29    Response = 2,
30    /// A frame of a stream (SPEC §3).
31    StreamFrame = 3,
32}
33
34/// `flags` bitmask: bits 0-1 are the [`InteractionShape`] (SPEC §2 field 3).
35pub const FLAG_SHAPE_MASK: u8 = 0b0000_0011;
36/// `flags` bit 2: the sealed bit — MUST be 1 for directed messages (SPEC §2 field 3).
37pub const FLAG_SEALED: u8 = 0b0000_0100;
38
39impl InteractionShape {
40    /// The shape encoded in a `flags` byte (bits 0-1). Unknown reserved values map to `None`.
41    #[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    /// This shape as its `flags` bits (0-1).
53    #[must_use]
54    pub fn as_bits(self) -> u8 {
55        self as u8
56    }
57}
58
59/// A stream frame's control header (SPEC §3), present iff the shape is `StreamFrame`. The state machine
60/// that drives these frames is WU4; the wire shape is final here.
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Streamable)]
62pub struct StreamHeader {
63    /// Frame kind (SPEC §3: OPEN=0, OPEN_ACK=1, DATA=2, CREDIT=3, CLOSE=4, CLOSE_ACK=5, RESET=6).
64    pub frame: u8,
65    /// Strictly-monotonic per-direction sequence number from 0 (the replay index, SPEC §3/§5.6).
66    pub seq: u64,
67    /// Credit-based flow-control window (SPEC §3 backpressure).
68    pub window: u32,
69}
70
71/// Stream frame kinds (SPEC §3). The wire field [`StreamHeader::frame`] is a `u8`; this enum names the
72/// values for WU4's state machine.
73#[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    /// The frame kind for a [`StreamHeader::frame`] byte, or `None` for an unknown/reserved value
86    /// (a malformed frame is rejected, never panics — SPEC §3 / §4 spirit).
87    #[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    /// This frame kind as its wire `u8` (the [`StreamHeader::frame`] value).
102    #[must_use]
103    pub fn as_u8(self) -> u8 {
104        self as u8
105    }
106}
107
108/// The e2e-sealed region (SPEC §5.2). WU1 defines the shape; WU2 fills `kem_enc` (the G1 ephemeral
109/// encapsulation) and `ciphertext` (`AEAD.Seal` of the [`InnerMessage`]).
110#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
111pub struct SealedPayload {
112    /// The DHKEM-over-G1 ephemeral encapsulation — a 48-byte compressed G1 point (SPEC §5.1). A
113    /// placeholder (zeros) until WU2 seals.
114    pub kem_enc: Bytes48,
115    /// `AEAD.Seal(key, nonce, aad=header, pt=InnerMessage)` (SPEC §5.2). Empty until WU2 seals.
116    pub ciphertext: Vec<u8>,
117}
118
119/// The sealed inner message (SPEC §5.2, FINAL field list — order is normative). Every field lives
120/// inside the AEAD-authenticated + signed region so a relay cannot tamper with any of it.
121///
122/// WU1 carries every field; the anti-replay (`counter`/`timestamp_ms`), expiry (`expires_at`), and
123/// signature (`sender_sig`) SEMANTICS are enforced by WU2/WU4 — the wire shape is final here.
124#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
125pub struct InnerMessage {
126    /// Re-bound message type, checked equal to the cleartext header (anti type-confusion, SPEC §5.2).
127    pub message_type: u32,
128    /// Re-bound correlation id, checked equal to the cleartext header (SPEC §5.2).
129    pub correlation_id: Bytes32,
130    /// Compression algorithm id of `payload` (SPEC §1.1).
131    pub compression: u8,
132    /// Declared original length of `payload` — the decompression-bomb bound (SPEC §1.1).
133    pub uncompressed_len: u32,
134    /// Per-(sender→recipient) strictly-monotonic anti-replay counter (SPEC §5.6). Enforced in WU4.
135    pub counter: u64,
136    /// Sender wall-clock Unix milliseconds — the freshness field (SPEC §5.6). Enforced in WU4.
137    pub timestamp_ms: u64,
138    /// Sender-controlled TTL, Unix milliseconds; 0 = no explicit expiry (SPEC §5.6b). Enforced in WU4.
139    pub expires_at: u64,
140    /// The COMPRESSED type-payload bytes (SPEC §1.1). The sole content region.
141    pub payload: Vec<u8>,
142    /// The mandatory 96-byte BLS G2 sender signature (SPEC §5.1). Zeros until WU2 signs.
143    pub sender_sig: Bytes96,
144}
145
146/// The base envelope (SPEC §2, field order normative). Fields 1-8 are the cleartext routing header;
147/// field 9 is the sealed region.
148#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
149pub struct DigMessageEnvelope {
150    /// Envelope format version (SPEC §2 field 1).
151    pub version: u8,
152    /// Cleartext message type (routing); re-bound inside the seal (SPEC §2 field 2).
153    pub message_type: u32,
154    /// Interaction-shape + sealed bitfield (SPEC §2 field 3).
155    pub flags: u8,
156    /// Random per initiating message; echoed by responses and every stream frame (SPEC §2 field 4).
157    pub correlation_id: Bytes32,
158    /// Sender DID launcher id (SPEC §2 field 5).
159    pub sender: Bytes32,
160    /// Recipient DID launcher id (SPEC §2 field 6).
161    pub recipient: Bytes32,
162    /// Sender key epoch for rotation disambiguation (SPEC §2 field 7).
163    pub sender_epoch: u32,
164    /// Present iff the shape is a stream frame (SPEC §2 field 8 / §3).
165    pub stream: Option<StreamHeader>,
166    /// The e2e-sealed region — all type-specific content (SPEC §2 field 9 / §5).
167    pub sealed: SealedPayload,
168}
169
170impl DigMessageEnvelope {
171    /// Serialize the cleartext header (fields 1-8, excluding the sealed region) — the bytes WU2 binds
172    /// as the AEAD AAD so an on-path party cannot alter routing metadata (SPEC §5.2).
173    ///
174    /// # Errors
175    /// [`MessageError::Codec`] if any field fails to serialize (should not happen for a well-formed
176    /// in-memory envelope).
177    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
192/// Serialize an envelope to its on-wire bytes, enforcing the [`MAX_ENVELOPE_BYTES`] cap (SPEC §1).
193///
194/// # Errors
195/// [`MessageError::EnvelopeTooLarge`] if the frame exceeds the cap; [`MessageError::Codec`] on a
196/// serialization failure.
197pub 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
210/// Decode an envelope from on-wire bytes, rejecting an over-cap frame BEFORE decoding and an unknown
211/// version after (SPEC §1, §2).
212///
213/// # Errors
214/// [`MessageError::EnvelopeTooLarge`] if the frame exceeds the cap; [`MessageError::Truncated`] on a
215/// short/malformed frame; [`MessageError::UnsupportedVersion`] for a newer version.
216pub 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    /// Build an envelope with deterministic, distinct field values for round-trip assertions.
237    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        // The header is a strict prefix of the full encoding (fields 1-8 precede field 9).
350        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}