Skip to main content

mkit_server/
upload.rs

1//! `UploadPack` stream framing (SPEC-TRANSPORT-CONNECT §6.1,
2//! SPEC-TRANSPORT §4.2).
3//!
4//! [`UploadValidator`] checks framing only: a 32-byte `pack_id`, a declared
5//! size within the binding's cap, chunks that repeat the header's `pack_id`
6//! at contiguous offsets without overrunning it, a bounded chunk count, and a
7//! `last` chunk at exactly the declared size. It never hashes: the pack sink
8//! verifies BLAKE3 before the pack becomes visible, and hashing here too
9//! would double the CPU cost of a 4 GiB upload.
10//!
11//! The canonical copy of `mkit serve`'s `UploadDrain`,
12//! `mkit-transport-connect` 0.4's `drain_upload` and the chunk loop of
13//! `vcs-worker`'s `upload_pack`. Each [`UploadError`] keeps today's text for
14//! both wire families: [`UploadError::ssh_message`] is `mkit serve`'s and
15//! [`UploadError::connect_message`] is `mkit-transport-connect`'s (the server
16//! it had until WP-M0-15, behind `mkit serve --http`). The old copies are
17//! gone: `mkit serve`'s in WP-M0-13, `vcs-worker`'s in WP-M0-17.
18
19use std::borrow::Cow;
20
21use mkit_core::protocol::PackKey;
22
23use crate::error::{Code, ServerError};
24
25/// Caps a binding applies to one upload. Each binding supplies its own:
26/// the native Connect server uses `PACK_BODY_LIMIT` and no chunk cap,
27/// `mkit serve` its per-connection byte and frame caps, and the Workers
28/// server its 64 MiB buffer cap.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct UploadLimits {
31    /// Largest `total_bytes` a header may declare.
32    pub max_total_bytes: u64,
33    /// Most chunks accepted before the `last` one, inclusive.
34    pub max_chunks: u32,
35}
36
37/// Where an upload stands after a chunk.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct Progress {
40    /// The `last` chunk arrived and the byte count matches the header.
41    pub complete: bool,
42}
43
44/// A completely framed upload.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct UploadDone {
47    /// The pack id the header declared.
48    pub key: PackKey,
49    /// Bytes received, equal to the declared `total_bytes`.
50    pub total: u64,
51}
52
53/// Why an upload stream was rejected.
54///
55/// The stream-shape variants ([`Self::HeaderMissing`],
56/// [`Self::UnexpectedMessage`]) and [`Self::DigestMismatch`] are raised by
57/// the binding, which decodes messages and owns the sink; the validator
58/// raises every other variant.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60#[non_exhaustive]
61pub enum UploadError {
62    /// The stream did not start with a header. `stream_empty` when it ended
63    /// before any message.
64    HeaderMissing {
65        /// No message arrived at all.
66        stream_empty: bool,
67    },
68    /// A message after the header was not a chunk. `header` when it was a
69    /// second header, otherwise a message with no body.
70    UnexpectedMessage {
71        /// The message was a second header.
72        header: bool,
73    },
74    /// A `pack_id` that is absent (`len: None`) or not 32 bytes long.
75    /// `chunk` when it was a chunk's rather than the header's.
76    BadPackId {
77        /// The id came from a chunk.
78        chunk: bool,
79        /// Its length, or `None` when absent.
80        len: Option<usize>,
81    },
82    /// A chunk's `pack_id` differs from the header's.
83    PackIdMismatch,
84    /// The header has no `total_bytes`.
85    TotalMissing,
86    /// The declared `total_bytes` is over the binding's cap.
87    TotalTooLarge {
88        /// Declared size.
89        total: u64,
90        /// The binding's cap.
91        cap: u64,
92    },
93    /// More chunks than the binding's cap arrived without `last`.
94    TooManyChunks,
95    /// A chunk has no `offset`.
96    OffsetMissing,
97    /// A chunk's `offset` is not the running byte count.
98    OffsetGap {
99        /// The chunk's offset.
100        offset: u64,
101        /// The running byte count.
102        expected: u64,
103    },
104    /// The running byte count overflowed `u64`.
105    ByteCountOverflow,
106    /// The chunks carry more bytes than the header declared.
107    Overrun,
108    /// A chunk arrived after the `last` one.
109    AfterLast,
110    /// The stream ended without a `last` chunk.
111    NoLast,
112    /// The `last` chunk arrived before the declared byte count.
113    LengthMismatch {
114        /// Bytes received.
115        received: u64,
116        /// Bytes the header declared.
117        declared: u64,
118    },
119    /// The received bytes do not hash to the header's `pack_id`.
120    DigestMismatch,
121}
122
123impl UploadError {
124    /// The Connect-side code: [`Code::ResourceExhausted`] for an oversized
125    /// declared total, [`Code::InvalidArgument`] for everything else.
126    #[must_use]
127    pub const fn code(self) -> Code {
128        match self {
129            Self::TotalTooLarge { .. } => Code::ResourceExhausted,
130            _ => Code::InvalidArgument,
131        }
132    }
133
134    /// `mkit serve`'s message, verbatim. Every one is sent with
135    /// `ERROR_CODE_INVALID_REQUEST`; the ssh binding maps codes itself.
136    #[must_use]
137    pub const fn ssh_message(self) -> &'static str {
138        match self {
139            Self::HeaderMissing { .. } => "PackChunk arrived without UploadPack header",
140            Self::UnexpectedMessage { .. } => "expected PackChunk after UploadPack",
141            Self::BadPackId { len: None, .. } => "pack_id missing",
142            Self::BadPackId { len: Some(_), .. } => "pack_id must be 32 bytes",
143            Self::PackIdMismatch => "PackChunk.pack_id does not match UploadPack",
144            Self::TotalMissing => "UploadPack.total_bytes is required",
145            Self::TotalTooLarge { .. } => "UploadPack.total_bytes exceeds server cap",
146            Self::TooManyChunks => "too many PackChunk frames before last=true",
147            Self::OffsetMissing => "PackChunk.offset is required",
148            Self::OffsetGap { .. } => "PackChunk.offset is not the expected next offset",
149            Self::ByteCountOverflow => "PackChunk byte count overflow",
150            Self::Overrun => "PackChunk data exceeds declared total_bytes",
151            // New: `mkit serve` stops reading at `last`, so it never sent this.
152            Self::AfterLast => "PackChunk after last=true",
153            // `mkit serve` reports a stream that ends early as a failed read.
154            Self::NoLast => "pack chunk read failed",
155            Self::LengthMismatch { .. } => "PackChunk stream ended before declared total_bytes",
156            Self::DigestMismatch => "uploaded pack bytes do not match UploadPack.pack_id",
157        }
158    }
159
160    /// `mkit-transport-connect`'s message, verbatim. `TotalMissing`,
161    /// `OffsetMissing`, `TooManyChunks` and `AfterLast` are new: that server
162    /// reads an absent `total_bytes` or `offset` as 0, has no chunk cap and
163    /// stops reading at `last`.
164    #[must_use]
165    pub fn connect_message(self) -> Cow<'static, str> {
166        Cow::Borrowed(match self {
167            Self::HeaderMissing { stream_empty: true } => "UploadPack: empty request stream",
168            Self::HeaderMissing {
169                stream_empty: false,
170            } => "UploadPack: first message MUST be `header`",
171            Self::UnexpectedMessage { header: true } => "UploadPack: saw a second `header` message",
172            Self::UnexpectedMessage { header: false } => {
173                "UploadPack: message with neither `header` nor `chunk` set"
174            }
175            Self::BadPackId { chunk: false, len } => {
176                return format!(
177                    "expected a 32-byte digest, got {} bytes",
178                    len.unwrap_or_default()
179                )
180                .into();
181            }
182            Self::BadPackId { chunk: true, .. } | Self::PackIdMismatch => {
183                "UploadPack: chunk.pack_id does not match header.pack_id"
184            }
185            Self::TotalMissing => "UploadPack: header.total_bytes is required",
186            Self::TotalTooLarge { total, cap } => {
187                return format!("UploadPack: total_bytes {total} exceeds the {cap}-byte cap")
188                    .into();
189            }
190            Self::TooManyChunks => {
191                "UploadPack: too many `chunk` messages before `chunk.last = true`"
192            }
193            Self::OffsetMissing => "UploadPack: chunk.offset is required",
194            Self::OffsetGap { offset, expected } => {
195                return format!(
196                    "UploadPack: chunk.offset {offset} does not match the expected offset {expected}"
197                )
198                .into();
199            }
200            Self::ByteCountOverflow | Self::Overrun => {
201                "UploadPack: received bytes exceed header.total_bytes"
202            }
203            Self::AfterLast => "UploadPack: message after `chunk.last = true`",
204            Self::NoLast => "UploadPack: stream ended without a `chunk.last = true` message",
205            Self::LengthMismatch { received, declared } => {
206                return format!(
207                    "UploadPack: received {received} bytes, header declared {declared}"
208                )
209                .into();
210            }
211            Self::DigestMismatch => {
212                "UploadPack: BLAKE3(received bytes) does not equal header.pack_id"
213            }
214        })
215    }
216}
217
218impl From<UploadError> for ServerError {
219    /// The Connect code and message.
220    fn from(err: UploadError) -> Self {
221        Self::new(err.code(), err.connect_message())
222    }
223}
224
225/// Validates one `UploadPack` stream's framing, chunk by chunk, without
226/// buffering or hashing. After any error the stream is dead: every later
227/// call returns that same error.
228#[derive(Debug, Clone)]
229pub struct UploadValidator {
230    key: PackKey,
231    declared: u64,
232    received: u64,
233    chunks: u32,
234    max_chunks: u32,
235    complete: bool,
236    failed: Option<UploadError>,
237}
238
239impl UploadValidator {
240    /// Start validating from the header's `pack_id` and `total_bytes`
241    /// (`None` when absent).
242    ///
243    /// # Errors
244    /// [`UploadError::BadPackId`], [`UploadError::TotalMissing`] or
245    /// [`UploadError::TotalTooLarge`].
246    pub fn new(
247        pack_id: Option<&[u8]>,
248        total_bytes: Option<u64>,
249        limits: UploadLimits,
250    ) -> Result<Self, UploadError> {
251        let key = pack_key(pack_id, false)?;
252        let declared = total_bytes.ok_or(UploadError::TotalMissing)?;
253        if declared > limits.max_total_bytes {
254            return Err(UploadError::TotalTooLarge {
255                total: declared,
256                cap: limits.max_total_bytes,
257            });
258        }
259        Ok(Self {
260            key,
261            declared,
262            received: 0,
263            chunks: 0,
264            max_chunks: limits.max_chunks,
265            complete: false,
266            failed: None,
267        })
268    }
269
270    /// Account for one chunk carrying `data_len` bytes.
271    ///
272    /// # Errors
273    /// [`UploadError::AfterLast`], [`UploadError::TooManyChunks`],
274    /// [`UploadError::BadPackId`], [`UploadError::PackIdMismatch`],
275    /// [`UploadError::OffsetMissing`], [`UploadError::OffsetGap`],
276    /// [`UploadError::ByteCountOverflow`], [`UploadError::Overrun`] or, on the
277    /// `last` chunk, [`UploadError::LengthMismatch`]. After any error, this
278    /// and every later call return that same error, and nothing the failed
279    /// chunk carried is counted.
280    pub fn push(
281        &mut self,
282        chunk_pack_id: Option<&[u8]>,
283        offset: Option<u64>,
284        data_len: usize,
285        last: bool,
286    ) -> Result<Progress, UploadError> {
287        if let Some(err) = self.failed {
288            return Err(err);
289        }
290        let result = self.accept(chunk_pack_id, offset, data_len, last);
291        if let Err(err) = result {
292            self.failed = Some(err);
293        }
294        result
295    }
296
297    fn accept(
298        &mut self,
299        chunk_pack_id: Option<&[u8]>,
300        offset: Option<u64>,
301        data_len: usize,
302        last: bool,
303    ) -> Result<Progress, UploadError> {
304        if self.complete {
305            return Err(UploadError::AfterLast);
306        }
307        self.chunks = self.chunks.saturating_add(1);
308        if self.chunks > self.max_chunks {
309            return Err(UploadError::TooManyChunks);
310        }
311        if pack_key(chunk_pack_id, true)? != self.key {
312            return Err(UploadError::PackIdMismatch);
313        }
314        let offset = offset.ok_or(UploadError::OffsetMissing)?;
315        if offset != self.received {
316            return Err(UploadError::OffsetGap {
317                offset,
318                expected: self.received,
319            });
320        }
321        let received = u64::try_from(data_len)
322            .ok()
323            .and_then(|len| self.received.checked_add(len))
324            .ok_or(UploadError::ByteCountOverflow)?;
325        if received > self.declared {
326            return Err(UploadError::Overrun);
327        }
328        if last && received != self.declared {
329            return Err(UploadError::LengthMismatch {
330                received,
331                declared: self.declared,
332            });
333        }
334        self.received = received;
335        self.complete = last;
336        Ok(Progress { complete: last })
337    }
338
339    /// End of stream.
340    ///
341    /// # Errors
342    /// The error that killed the stream, if any; otherwise
343    /// [`UploadError::NoLast`] if the `last` chunk never arrived.
344    pub fn finish(self) -> Result<UploadDone, UploadError> {
345        if let Some(err) = self.failed {
346            return Err(err);
347        }
348        if !self.complete {
349            return Err(UploadError::NoLast);
350        }
351        Ok(UploadDone {
352            key: self.key,
353            total: self.received,
354        })
355    }
356
357    /// The pack id the header declared.
358    #[must_use]
359    pub const fn key(&self) -> PackKey {
360        self.key
361    }
362
363    /// The `total_bytes` the header declared.
364    #[must_use]
365    pub const fn declared(&self) -> u64 {
366        self.declared
367    }
368
369    /// Bytes accepted so far.
370    #[must_use]
371    pub const fn received(&self) -> u64 {
372        self.received
373    }
374}
375
376fn pack_key(id: Option<&[u8]>, chunk: bool) -> Result<PackKey, UploadError> {
377    let id = id.ok_or(UploadError::BadPackId { chunk, len: None })?;
378    <[u8; 32]>::try_from(id)
379        .map(PackKey::new)
380        .map_err(|_| UploadError::BadPackId {
381            chunk,
382            len: Some(id.len()),
383        })
384}
385
386#[cfg(test)]
387mod tests;
388
389pub(crate) mod marker;
390pub(crate) mod receipt;
391pub(crate) mod ticket_auth;
392/// Stateless authenticated upload tickets.
393pub mod token;