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
84/// The e2e-sealed region (SPEC §5.2). WU1 defines the shape; WU2 fills `kem_enc` (the G1 ephemeral
85/// encapsulation) and `ciphertext` (`AEAD.Seal` of the [`InnerMessage`]).
86#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
87pub struct SealedPayload {
88    /// The DHKEM-over-G1 ephemeral encapsulation — a 48-byte compressed G1 point (SPEC §5.1). A
89    /// placeholder (zeros) until WU2 seals.
90    pub kem_enc: Bytes48,
91    /// `AEAD.Seal(key, nonce, aad=header, pt=InnerMessage)` (SPEC §5.2). Empty until WU2 seals.
92    pub ciphertext: Vec<u8>,
93}
94
95/// The sealed inner message (SPEC §5.2, FINAL field list — order is normative). Every field lives
96/// inside the AEAD-authenticated + signed region so a relay cannot tamper with any of it.
97///
98/// WU1 carries every field; the anti-replay (`counter`/`timestamp_ms`), expiry (`expires_at`), and
99/// signature (`sender_sig`) SEMANTICS are enforced by WU2/WU4 — the wire shape is final here.
100#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
101pub struct InnerMessage {
102    /// Re-bound message type, checked equal to the cleartext header (anti type-confusion, SPEC §5.2).
103    pub message_type: u32,
104    /// Re-bound correlation id, checked equal to the cleartext header (SPEC §5.2).
105    pub correlation_id: Bytes32,
106    /// Compression algorithm id of `payload` (SPEC §1.1).
107    pub compression: u8,
108    /// Declared original length of `payload` — the decompression-bomb bound (SPEC §1.1).
109    pub uncompressed_len: u32,
110    /// Per-(sender→recipient) strictly-monotonic anti-replay counter (SPEC §5.6). Enforced in WU4.
111    pub counter: u64,
112    /// Sender wall-clock Unix milliseconds — the freshness field (SPEC §5.6). Enforced in WU4.
113    pub timestamp_ms: u64,
114    /// Sender-controlled TTL, Unix milliseconds; 0 = no explicit expiry (SPEC §5.6b). Enforced in WU4.
115    pub expires_at: u64,
116    /// The COMPRESSED type-payload bytes (SPEC §1.1). The sole content region.
117    pub payload: Vec<u8>,
118    /// The mandatory 96-byte BLS G2 sender signature (SPEC §5.1). Zeros until WU2 signs.
119    pub sender_sig: Bytes96,
120}
121
122/// The base envelope (SPEC §2, field order normative). Fields 1-8 are the cleartext routing header;
123/// field 9 is the sealed region.
124#[derive(Debug, Clone, PartialEq, Eq, Streamable)]
125pub struct DigMessageEnvelope {
126    /// Envelope format version (SPEC §2 field 1).
127    pub version: u8,
128    /// Cleartext message type (routing); re-bound inside the seal (SPEC §2 field 2).
129    pub message_type: u32,
130    /// Interaction-shape + sealed bitfield (SPEC §2 field 3).
131    pub flags: u8,
132    /// Random per initiating message; echoed by responses and every stream frame (SPEC §2 field 4).
133    pub correlation_id: Bytes32,
134    /// Sender DID launcher id (SPEC §2 field 5).
135    pub sender: Bytes32,
136    /// Recipient DID launcher id (SPEC §2 field 6).
137    pub recipient: Bytes32,
138    /// Sender key epoch for rotation disambiguation (SPEC §2 field 7).
139    pub sender_epoch: u32,
140    /// Present iff the shape is a stream frame (SPEC §2 field 8 / §3).
141    pub stream: Option<StreamHeader>,
142    /// The e2e-sealed region — all type-specific content (SPEC §2 field 9 / §5).
143    pub sealed: SealedPayload,
144}
145
146impl DigMessageEnvelope {
147    /// Serialize the cleartext header (fields 1-8, excluding the sealed region) — the bytes WU2 binds
148    /// as the AEAD AAD so an on-path party cannot alter routing metadata (SPEC §5.2).
149    ///
150    /// # Errors
151    /// [`MessageError::Codec`] if any field fails to serialize (should not happen for a well-formed
152    /// in-memory envelope).
153    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
168/// Serialize an envelope to its on-wire bytes, enforcing the [`MAX_ENVELOPE_BYTES`] cap (SPEC §1).
169///
170/// # Errors
171/// [`MessageError::EnvelopeTooLarge`] if the frame exceeds the cap; [`MessageError::Codec`] on a
172/// serialization failure.
173pub 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
186/// Decode an envelope from on-wire bytes, rejecting an over-cap frame BEFORE decoding and an unknown
187/// version after (SPEC §1, §2).
188///
189/// # Errors
190/// [`MessageError::EnvelopeTooLarge`] if the frame exceeds the cap; [`MessageError::Truncated`] on a
191/// short/malformed frame; [`MessageError::UnsupportedVersion`] for a newer version.
192pub 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    /// Build an envelope with deterministic, distinct field values for round-trip assertions.
213    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        // The header is a strict prefix of the full encoding (fields 1-8 precede field 9).
326        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}