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}