moq_secure/
wire.rs

1use crate::crypto::{aead_decrypt, aead_encrypt, sha256_digest};
2use crate::error::MoqSecureError;
3use crate::key_store::KeyStore;
4use ed25519_dalek::{Signer, Verifier};
5
6pub const MAGIC: [u8; 4] = *b"MOQS";
7pub const VERSION: u8 = 1;
8
9pub const SIG_SLOT_LEN: usize = 64;
10pub const AEAD_TAG_LEN: usize = 16;
11pub const PAD_LEN_FIELD_LEN: usize = 4;
12
13// magic(4) | version(1) | key_id(1) | ctr(8) |
14// n_signed(1) | sig_flag(1) | encrypted(1)
15pub const FIXED_HEADER_LEN: usize = 4 + 1 + 1 + 8 + 1 + 1 + 1;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct WireHeader {
19    pub magic: [u8; 4],
20    pub version: u8,
21    pub key_id: u8,
22    pub ctr: u64,
23    pub n_signed: u8,
24    pub sig_flag: u8,
25    pub encrypted: u8,
26}
27
28impl WireHeader {
29    pub fn encode(&self) -> Vec<u8> {
30        let mut v = Vec::with_capacity(FIXED_HEADER_LEN);
31
32        v.extend_from_slice(&self.magic);
33        v.push(self.version);
34        v.push(self.key_id);
35        v.extend_from_slice(&self.ctr.to_be_bytes());
36        v.push(self.n_signed);
37        v.push(self.sig_flag);
38        v.push(self.encrypted);
39
40        v
41    }
42
43    pub fn aad(&self) -> Vec<u8> {
44        self.encode()
45    }
46
47    pub fn validate(&self) -> Result<(), MoqSecureError> {
48        if self.magic != MAGIC {
49            return Err(MoqSecureError::InvalidMagic);
50        }
51
52        if self.version != VERSION {
53            return Err(MoqSecureError::UnsupportedVersion(self.version));
54        }
55
56        if self.sig_flag != 0 && self.sig_flag != 1 {
57            return Err(MoqSecureError::InvalidSigFlag(self.sig_flag));
58        }
59
60        if self.encrypted != 0 && self.encrypted != 1 {
61            return Err(MoqSecureError::InvalidEncryptedFlag(self.encrypted));
62        }
63
64        if self.n_signed == 0 && self.sig_flag != 0 {
65            return Err(MoqSecureError::SigningMismatch);
66        }
67
68        Ok(())
69    }
70}
71
72#[derive(Debug, Clone)]
73pub struct Frame {
74    pub header: WireHeader,
75
76    // For encrypted frames, payload is ciphertext without the AEAD tag.
77    // For unencrypted frames, payload is:
78    // padLen(4-byte big endian) || padding || plaintext.
79    pub payload: Vec<u8>,
80
81    pub tag: [u8; AEAD_TAG_LEN],
82    pub signature: Option<[u8; SIG_SLOT_LEN]>,
83}
84
85impl Frame {
86    pub fn parse(frame: &[u8]) -> Result<Self, MoqSecureError> {
87        if frame.len() < FIXED_HEADER_LEN {
88            return Err(MoqSecureError::TruncatedFrame);
89        }
90
91        let (h_bytes, rest0) = frame.split_at(FIXED_HEADER_LEN);
92        let mut idx = 0;
93
94        let mut magic = [0u8; 4];
95        magic.copy_from_slice(&h_bytes[idx..idx + 4]);
96        idx += 4;
97
98        let version = h_bytes[idx];
99        idx += 1;
100
101        let key_id = h_bytes[idx];
102        idx += 1;
103
104        let mut ctr_bytes = [0u8; 8];
105        ctr_bytes.copy_from_slice(&h_bytes[idx..idx + 8]);
106        idx += 8;
107
108        let ctr = u64::from_be_bytes(ctr_bytes);
109
110        let n_signed = h_bytes[idx];
111        idx += 1;
112
113        let sig_flag = h_bytes[idx];
114        idx += 1;
115
116        let encrypted = h_bytes[idx];
117
118        let header = WireHeader {
119            magic,
120            version,
121            key_id,
122            ctr,
123            n_signed,
124            sig_flag,
125            encrypted,
126        };
127
128        header.validate()?;
129
130        let sig_len = if header.sig_flag == 1 {
131            SIG_SLOT_LEN
132        } else {
133            0
134        };
135
136        if rest0.len() < sig_len {
137            return Err(MoqSecureError::TruncatedFrame);
138        }
139
140        let (payload_and_tag, signature_bytes) = if sig_len == 0 {
141            (rest0, None)
142        } else {
143            let split_at = rest0.len() - SIG_SLOT_LEN;
144            let (payload, signature) = rest0.split_at(split_at);
145            (payload, Some(signature))
146        };
147
148        let signature = if let Some(sig_bytes) = signature_bytes {
149            if sig_bytes.len() != SIG_SLOT_LEN {
150                return Err(MoqSecureError::TruncatedFrame);
151            }
152
153            let mut sig = [0u8; SIG_SLOT_LEN];
154            sig.copy_from_slice(sig_bytes);
155
156            if sig == [0u8; SIG_SLOT_LEN] {
157                return Err(MoqSecureError::InvalidSignature);
158            }
159
160            Some(sig)
161        } else {
162            None
163        };
164
165        if header.encrypted == 1 {
166            if payload_and_tag.len() < AEAD_TAG_LEN {
167                return Err(MoqSecureError::CiphertextTooShort);
168            }
169
170            let split_at = payload_and_tag.len() - AEAD_TAG_LEN;
171            let ciphertext = &payload_and_tag[..split_at];
172            let tag_bytes = &payload_and_tag[split_at..];
173
174            let tag: [u8; AEAD_TAG_LEN] = tag_bytes
175                .try_into()
176                .map_err(|_| MoqSecureError::CiphertextTooShort)?;
177
178            Ok(Self {
179                header,
180                payload: ciphertext.to_vec(),
181                tag,
182                signature,
183            })
184        } else {
185            Ok(Self {
186                header,
187                payload: payload_and_tag.to_vec(),
188                tag: [0u8; AEAD_TAG_LEN],
189                signature,
190            })
191        }
192    }
193
194    pub fn serialize(&self) -> Vec<u8> {
195        let mut out = Vec::new();
196
197        out.extend_from_slice(&self.header.encode());
198
199        if self.header.encrypted == 1 {
200            out.extend_from_slice(&self.payload);
201            out.extend_from_slice(&self.tag);
202        } else {
203            out.extend_from_slice(&self.payload);
204        }
205
206        if self.header.sig_flag == 1 {
207            if let Some(signature) = self.signature {
208                out.extend_from_slice(&signature);
209            } else {
210                out.extend_from_slice(&[0u8; SIG_SLOT_LEN]);
211            }
212        }
213
214        out
215    }
216
217    pub fn aad_bytes(&self) -> Vec<u8> {
218        self.header.aad()
219    }
220
221    pub fn digest_for_signature(&self) -> [u8; 32] {
222        let header_bytes = self.header.encode();
223
224        let mut data = Vec::with_capacity(
225            header_bytes.len()
226                + self.payload.len()
227                + if self.header.encrypted == 1 {
228                    AEAD_TAG_LEN
229                } else {
230                    0
231                },
232        );
233
234        data.extend_from_slice(&header_bytes);
235        data.extend_from_slice(&self.payload);
236
237        if self.header.encrypted == 1 {
238            data.extend_from_slice(&self.tag);
239        }
240
241        sha256_digest(&data)
242    }
243
244    pub fn decode_plaintext_with_key_store(
245        &self,
246        key_store: &dyn KeyStore,
247        broadcaster_public_key: &ed25519_dalek::VerifyingKey,
248        lease_remaining: &mut u8,
249    ) -> Result<Vec<u8>, MoqSecureError> {
250        if self.header.n_signed == 0 {
251            if self.header.sig_flag != 0 || self.signature.is_some() {
252                return Err(MoqSecureError::SigningMismatch);
253            }
254        } else if self.header.sig_flag == 1 {
255            let signature_bytes = self
256                .signature
257                .ok_or(MoqSecureError::InvalidSignature)?;
258
259            let signature =
260                ed25519_dalek::Signature::from_bytes(&signature_bytes);
261
262            let digest = self.digest_for_signature();
263
264            broadcaster_public_key
265                .verify(&digest, &signature)
266                .map_err(|_| MoqSecureError::InvalidSignature)?;
267
268            *lease_remaining = self.header.n_signed;
269        } else {
270            if *lease_remaining == 0 {
271                return Err(MoqSecureError::InvalidSignature);
272            }
273
274            *lease_remaining -= 1;
275        }
276
277        let padded_plaintext = if self.header.encrypted == 1 {
278            let key = key_store
279                .aead_key(self.header.key_id)
280                .ok_or(MoqSecureError::InvalidKeyId(self.header.key_id))?;
281
282            aead_decrypt(
283                key,
284                self.header.key_id,
285                self.header.ctr,
286                &self.aad_bytes(),
287                &self.payload,
288                &self.tag,
289            )?
290        } else {
291            self.payload.clone()
292        };
293
294        if padded_plaintext.len() < PAD_LEN_FIELD_LEN {
295            return Err(MoqSecureError::InvalidPadLength);
296        }
297
298        let mut pad_len_bytes = [0u8; PAD_LEN_FIELD_LEN];
299        pad_len_bytes.copy_from_slice(&padded_plaintext[..PAD_LEN_FIELD_LEN]);
300
301        let pad_len = u32::from_be_bytes(pad_len_bytes) as usize;
302        let content_start = PAD_LEN_FIELD_LEN
303            .checked_add(pad_len)
304            .ok_or(MoqSecureError::InvalidPadLength)?;
305
306        if content_start > padded_plaintext.len() {
307            return Err(MoqSecureError::InvalidPadLength);
308        }
309
310        Ok(padded_plaintext[content_start..].to_vec())
311    }
312}
313
314pub fn encrypt_frame(
315    key_store: &dyn KeyStore,
316    broadcaster_private_key: &ed25519_dalek::SigningKey,
317    key_id: u8,
318    ctr: u64,
319    n_signed: u8,
320    maybe_sign: bool,
321    encrypted: u8,
322    pad_len: u32,
323    plaintext: &[u8],
324) -> Result<Frame, MoqSecureError> {
325    if encrypted != 0 && encrypted != 1 {
326        return Err(MoqSecureError::InvalidEncryptedFlag(encrypted));
327    }
328
329    let sig_flag = if n_signed != 0 && maybe_sign { 1 } else { 0 };
330
331    let header = WireHeader {
332        magic: MAGIC,
333        version: VERSION,
334        key_id,
335        ctr,
336        n_signed,
337        sig_flag,
338        encrypted,
339    };
340
341    header.validate()?;
342
343    let pad_len_usize = pad_len as usize;
344    let mut padded_plaintext = Vec::with_capacity(
345        PAD_LEN_FIELD_LEN + pad_len_usize + plaintext.len(),
346    );
347
348    padded_plaintext.extend_from_slice(&pad_len.to_be_bytes());
349    padded_plaintext.resize(PAD_LEN_FIELD_LEN + pad_len_usize, 0);
350    padded_plaintext.extend_from_slice(plaintext);
351
352    let frame_without_signature = if encrypted == 1 {
353        let key = key_store
354            .aead_key(key_id)
355            .ok_or(MoqSecureError::InvalidKeyId(key_id))?;
356
357        let (ciphertext, tag) = aead_encrypt(
358            key,
359            key_id,
360            ctr,
361            &header.aad(),
362            &padded_plaintext,
363        );
364
365        Frame {
366            header,
367            payload: ciphertext,
368            tag,
369            signature: None,
370        }
371    } else {
372        Frame {
373            header,
374            payload: padded_plaintext,
375            tag: [0u8; AEAD_TAG_LEN],
376            signature: None,
377        }
378    };
379
380    if sig_flag == 1 {
381        let digest = frame_without_signature.digest_for_signature();
382        let signature = broadcaster_private_key.sign(&digest);
383
384        Ok(Frame {
385            signature: Some(signature.to_bytes()),
386            ..frame_without_signature
387        })
388    } else {
389        Ok(frame_without_signature)
390    }
391}
392
393pub fn decrypt_frame(
394    key_store: &dyn KeyStore,
395    broadcaster_public_key: &ed25519_dalek::VerifyingKey,
396    lease_remaining: &mut u8,
397    frame_bytes: &[u8],
398) -> Result<Vec<u8>, MoqSecureError> {
399    let frame = Frame::parse(frame_bytes)?;
400
401    frame.decode_plaintext_with_key_store(
402        key_store,
403        broadcaster_public_key,
404        lease_remaining,
405    )
406}
moq_secure/wire.rs

moq_secure/
wire.rs
