Skip to main content

mkit_server/ssh/
budget.rs

1//! Per-connection resource caps (SPEC-TRANSPORT §4.4).
2//!
3//! A session is driven by one remote client: an ssh forced command or an
4//! enc peer. Bounding its cumulative work stops a misbehaving client from
5//! pinning the serving process. The numbers are the ones SPEC-TRANSPORT
6//! §4.4 cites.
7
8use mkit_rpc::mkit::rpc::v1::ssh::{SshFrame, ssh_frame};
9
10use crate::upload::UploadLimits;
11
12/// Most top-level frames a session reads after `Hello`. The chunk frames an
13/// upload reads are not counted here; [`UploadLimits::max_chunks`] caps
14/// them instead.
15pub const MAX_FRAMES_PER_CONN: u32 = 10_000;
16
17/// Most estimated request bytes per session (1 GiB); see
18/// [`frame_byte_estimate`]. It also caps an upload's declared size.
19pub const MAX_BYTES_PER_CONN: u64 = 1024 * 1024 * 1024;
20
21/// The upload caps of the ssh wire: a declared size of at most
22/// [`MAX_BYTES_PER_CONN`] and at most [`MAX_FRAMES_PER_CONN`] chunks. A
23/// binding builds its pipeline's `PipelineConfig` with these.
24#[must_use]
25pub const fn upload_limits() -> UploadLimits {
26    UploadLimits {
27        max_total_bytes: MAX_BYTES_PER_CONN,
28        max_chunks: MAX_FRAMES_PER_CONN,
29    }
30}
31
32/// A frame's cost against [`MAX_BYTES_PER_CONN`], without re-encoding: a
33/// chunk's data length, the `total_bytes` a header declares, or 64 for a
34/// small control frame. An upload is charged its declared size once, by its
35/// header; the chunks it then reads are not charged again.
36#[must_use]
37pub fn frame_byte_estimate(f: &SshFrame) -> u64 {
38    use ssh_frame::Body;
39    match &f.body {
40        Some(Body::PackChunk(c)) => c.data.as_ref().map_or(0, Vec::len) as u64,
41        Some(Body::UploadPack(h)) => h.total_bytes.unwrap_or(0),
42        Some(Body::DownloadPackHeader(h)) => h.total_bytes.unwrap_or(0),
43        _ => 64,
44    }
45}
46
47/// The running totals of one session.
48#[derive(Debug, Default)]
49pub(super) struct Budget {
50    frames: u32,
51    bytes: u64,
52}
53
54impl Budget {
55    /// Charge one top-level frame.
56    ///
57    /// # Errors
58    /// The message of the cap it exceeds.
59    pub(super) fn charge(&mut self, frame: &SshFrame) -> Result<(), &'static str> {
60        self.frames = self.frames.saturating_add(1);
61        if self.frames > MAX_FRAMES_PER_CONN {
62            return Err("per-connection frame budget exceeded");
63        }
64        self.bytes = self.bytes.saturating_add(frame_byte_estimate(frame));
65        if self.bytes > MAX_BYTES_PER_CONN {
66            return Err("per-connection byte budget exceeded");
67        }
68        Ok(())
69    }
70}