1use std::borrow::Cow;
20
21use mkit_core::protocol::PackKey;
22
23use crate::error::{Code, ServerError};
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct UploadLimits {
31 pub max_total_bytes: u64,
33 pub max_chunks: u32,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct Progress {
40 pub complete: bool,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct UploadDone {
47 pub key: PackKey,
49 pub total: u64,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60#[non_exhaustive]
61pub enum UploadError {
62 HeaderMissing {
65 stream_empty: bool,
67 },
68 UnexpectedMessage {
71 header: bool,
73 },
74 BadPackId {
77 chunk: bool,
79 len: Option<usize>,
81 },
82 PackIdMismatch,
84 TotalMissing,
86 TotalTooLarge {
88 total: u64,
90 cap: u64,
92 },
93 TooManyChunks,
95 OffsetMissing,
97 OffsetGap {
99 offset: u64,
101 expected: u64,
103 },
104 ByteCountOverflow,
106 Overrun,
108 AfterLast,
110 NoLast,
112 LengthMismatch {
114 received: u64,
116 declared: u64,
118 },
119 DigestMismatch,
121}
122
123impl UploadError {
124 #[must_use]
127 pub const fn code(self) -> Code {
128 match self {
129 Self::TotalTooLarge { .. } => Code::ResourceExhausted,
130 _ => Code::InvalidArgument,
131 }
132 }
133
134 #[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 Self::AfterLast => "PackChunk after last=true",
153 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 #[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 fn from(err: UploadError) -> Self {
221 Self::new(err.code(), err.connect_message())
222 }
223}
224
225#[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 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 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 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 #[must_use]
359 pub const fn key(&self) -> PackKey {
360 self.key
361 }
362
363 #[must_use]
365 pub const fn declared(&self) -> u64 {
366 self.declared
367 }
368
369 #[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;
392pub mod token;