krabka_kraft_core/
snapshot_fetch.rs

1//! Pure reassembly state machine for a follower that fetches a KIP-630
2//! snapshot over `FetchSnapshot` (api key 59).
3//!
4//! This module is IO-free. The async engine (`controller.rs`) owns the
5//! transport and applies the [`SnapshotFetchStep`] that this module returns.
6//!
7//! Chunks must arrive in order, that is, `position` must equal the number of
8//! bytes received so far. Any mismatch or a changed snapshot id aborts the
9//! transfer, so the engine restarts cleanly against the current leader.
10
11use {
12    crate::types::NodeId,
13    bytes::{Bytes, BytesMut},
14    krabka_macros::RefinedNewtype,
15    krabka_units::prelude::{ByteSize, ByteSizeExt as _, gibibytes},
16    krabka_verified::{SnapshotChunkDecision, snapshot_chunk_admission},
17    refined_type::rule::MinMaxU64,
18};
19
20/// Snapshot identity: (`end_offset` exclusive, epoch), as KIP-630 `SnapshotId`.
21pub type SnapshotId = (i64, i32);
22
23/// Hard ceiling on the total bytes a follower will reassemble for one snapshot.
24///
25/// The size of the cluster metadata bounds the `__cluster_metadata` snapshot:
26/// topics, partitions, ACLs, configs, and broker registrations. That is tens of
27/// MiB even for very large clusters. 1 GiB is generous but finite. No
28/// legitimate cluster reaches it, and it caps the memory that a peer which the
29/// follower believes is the leader can force it to allocate (denial-of-service
30/// finding M-3). A peer that declares or streams more than this aborts the
31/// transfer.
32pub const METADATA_SNAPSHOT_FETCH_HARD_MAX: ByteSize = gibibytes(1);
33
34const METADATA_SNAPSHOT_FETCH_HARD_MAX_BYTES: u64 = 1_073_741_824;
35
36/// Validated deployment limit for one metadata snapshot fetch.
37///
38/// Operators may lower this limit. No configuration can raise the fixed 1 GiB
39/// security ceiling.
40#[derive(Clone, Copy, Debug, Eq, PartialEq, RefinedNewtype)]
41#[refined(
42    rule(MinMaxU64<1, METADATA_SNAPSHOT_FETCH_HARD_MAX_BYTES>),
43    quantity = ByteSize,
44    label = "metadata snapshot fetch max",
45    quantity_getter = byte_size,
46    default = METADATA_SNAPSHOT_FETCH_HARD_MAX
47)]
48pub struct MetadataSnapshotFetchMax(u64);
49
50/// In-flight reassembly of one snapshot from one leader.
51#[derive(Debug)]
52pub struct SnapshotFetchState {
53    pub snapshot_id: SnapshotId,
54    pub leader_id: NodeId,
55    buf: BytesMut,
56    size: Option<i64>,
57    max_size: MetadataSnapshotFetchMax,
58}
59
60/// What the engine should do after it feeds a chunk in.
61#[derive(Debug, PartialEq, Eq)]
62pub enum SnapshotFetchStep {
63    /// Request the next byte range that starts at `next_position`.
64    Continue { next_position: i64 },
65    /// All bytes arrived; this variant holds the assembled snapshot.
66    Complete(Bytes),
67    /// Abort after an id mismatch, an out-of-order chunk, or a leader change.
68    /// The engine discards this state and falls back to a normal Fetch.
69    Restart,
70}
71
72impl SnapshotFetchState {
73    #[must_use]
74    pub fn new(snapshot_id: SnapshotId, leader_id: NodeId) -> Self {
75        Self::with_max(snapshot_id, leader_id, MetadataSnapshotFetchMax::default())
76    }
77
78    /// Start reassembly with a validated deployment-specific maximum.
79    #[must_use]
80    pub fn with_max(
81        snapshot_id: SnapshotId,
82        leader_id: NodeId,
83        max_size: MetadataSnapshotFetchMax,
84    ) -> Self {
85        Self {
86            snapshot_id,
87            leader_id,
88            buf: BytesMut::new(),
89            size: None,
90            max_size,
91        }
92    }
93
94    /// The byte position to request next, that is, the bytes received so far.
95    #[must_use]
96    pub fn next_position(&self) -> i64 { i64::try_from(self.buf.len()).unwrap_or(i64::MAX) }
97
98    /// Feed one response chunk.
99    ///
100    /// `id`, `size`, and `position` come from the `FetchSnapshot` response.
101    /// `chunk` holds the `unaligned_records` bytes.
102    pub fn on_chunk(
103        &mut self,
104        id: SnapshotId,
105        size: i64,
106        position: i64,
107        chunk: &[u8],
108    ) -> SnapshotFetchStep {
109        let Ok(received) = i64::try_from(self.buf.len()) else {
110            return SnapshotFetchStep::Restart;
111        };
112        let Ok(chunk_len) = i64::try_from(chunk.len()) else {
113            return SnapshotFetchStep::Restart;
114        };
115        let decision = snapshot_chunk_admission(
116            id == self.snapshot_id,
117            self.size,
118            received,
119            size,
120            position,
121            chunk_len,
122            self.max_size.byte_size().bytes_i64(),
123        );
124        match decision {
125            SnapshotChunkDecision::Restart => SnapshotFetchStep::Restart,
126            SnapshotChunkDecision::Continue { next_position } => {
127                self.size = Some(size);
128                self.buf.extend_from_slice(chunk);
129                SnapshotFetchStep::Continue { next_position }
130            }
131            SnapshotChunkDecision::Complete => {
132                self.size = Some(size);
133                self.buf.extend_from_slice(chunk);
134                SnapshotFetchStep::Complete(self.buf.split().freeze())
135            }
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use {super::*, assert2::assert, krabka_units::prelude::bytes};
143
144    fn start_three_byte_chunk(s: &mut SnapshotFetchState) {
145        assert2::assert!(
146            s.on_chunk((10, 1), 6, 0, b"abc") == SnapshotFetchStep::Continue { next_position: 3 }
147        );
148    }
149
150    #[test]
151    fn snapshot_fetch_max_accepts_only_whole_bytes_within_the_security_ceiling() {
152        assert2::assert!(MetadataSnapshotFetchMax::new(bytes(1)).is_ok());
153        assert2::assert!(MetadataSnapshotFetchMax::new(gibibytes(1)).is_ok());
154        assert2::assert!(MetadataSnapshotFetchMax::new(ByteSize::from_bytes_f64(0.5)).is_err());
155        assert2::assert!(MetadataSnapshotFetchMax::new(bytes(0)).is_err());
156        assert2::assert!(MetadataSnapshotFetchMax::new(bytes(1_073_741_825)).is_err());
157    }
158
159    #[test]
160    fn configured_snapshot_fetch_max_may_lower_but_not_raise_the_ceiling() {
161        let max = MetadataSnapshotFetchMax::new(bytes(3)).expect("valid configured maximum");
162        let mut oversized = SnapshotFetchState::with_max((10, 1), NodeId(2), max);
163        assert2::assert!(oversized.on_chunk((10, 1), 4, 0, b"abc") == SnapshotFetchStep::Restart);
164
165        let mut exact = SnapshotFetchState::with_max((10, 1), NodeId(2), max);
166        assert2::assert!(
167            exact.on_chunk((10, 1), 3, 0, b"abc")
168                == SnapshotFetchStep::Complete(Bytes::from_static(b"abc"))
169        );
170    }
171
172    /// A declared total the receiver will not accept: negative, changed
173    /// mid-transfer, or contradicted by a chunk that runs past it.
174    ///
175    /// Each is its own guard and each is silent unless something asks for it.
176    /// A leader that could shrink `size` after the fact, or overrun it, would
177    /// push `buf` past the bound the cap is supposed to give.
178    #[test]
179    fn a_declared_size_that_cannot_be_trusted_restarts_the_transfer() {
180        // A negative total, rejected on its own rather than only alongside a
181        // wrong id or position.
182        let mut negative = SnapshotFetchState::new((10, 1), NodeId(2));
183        assert2::assert!(negative.on_chunk((10, 1), -1, 0, b"") == SnapshotFetchStep::Restart);
184
185        // Zero is a real total: an empty snapshot completes immediately.
186        let mut empty = SnapshotFetchState::new((10, 1), NodeId(2));
187        assert2::assert!(
188            empty.on_chunk((10, 1), 0, 0, b"") == SnapshotFetchStep::Complete(Bytes::new())
189        );
190
191        // The total may not change once declared.
192        let mut changed = SnapshotFetchState::new((10, 1), NodeId(2));
193        assert2::assert!(
194            changed.on_chunk((10, 1), 6, 0, b"abc")
195                == SnapshotFetchStep::Continue { next_position: 3 }
196        );
197        assert2::assert!(changed.on_chunk((10, 1), 9, 3, b"def") == SnapshotFetchStep::Restart);
198
199        // A chunk may not carry more than the remaining bytes.
200        let mut overshoot = SnapshotFetchState::new((10, 1), NodeId(2));
201        assert2::assert!(
202            overshoot.on_chunk((10, 1), 6, 0, b"abc")
203                == SnapshotFetchStep::Continue { next_position: 3 }
204        );
205        assert2::assert!(overshoot.on_chunk((10, 1), 6, 3, b"defgh") == SnapshotFetchStep::Restart);
206    }
207
208    #[test]
209    fn assembles_in_order_chunks_to_complete() {
210        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
211        assert2::assert!(s.next_position() == 0);
212        let step = s.on_chunk((10, 1), 6, 0, b"abc");
213        assert2::assert!(step == SnapshotFetchStep::Continue { next_position: 3 });
214        let step = s.on_chunk((10, 1), 6, 3, b"def");
215        match step {
216            SnapshotFetchStep::Complete(b) => assert2::assert!(b.as_ref() == b"abcdef"),
217            other => panic!("expected Complete, got {other:?}"),
218        }
219    }
220
221    #[test]
222    fn out_of_order_position_restarts() {
223        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
224        let _ = s.on_chunk((10, 1), 6, 0, b"abc");
225        assert2::assert!(s.on_chunk((10, 1), 6, 99, b"def") == SnapshotFetchStep::Restart);
226    }
227
228    #[test]
229    fn mismatched_id_restarts() {
230        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
231        assert2::assert!(s.on_chunk((11, 1), 6, 0, b"abc") == SnapshotFetchStep::Restart);
232    }
233
234    #[test]
235    fn chunk_overshooting_declared_size_restarts() {
236        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
237        // Leader declares size 3 but streams 5 bytes in one chunk.
238        let step = s.on_chunk((10, 1), 3, 0, b"abcde");
239        // buf must not have been blown past the declared size.
240        assert2::assert!(step == SnapshotFetchStep::Restart);
241        assert2::assert!(s.next_position() == 0);
242    }
243
244    #[yare::parameterized(
245        duplicate_chunk = { (10, 1), 6, 0, b"abc".as_slice() },
246        overlapping_chunk = { (10, 1), 6, 2, b"cde".as_slice() },
247        starts_after_gap = { (10, 1), 6, 4, b"ef".as_slice() },
248        changed_snapshot_id = { (11, 1), 6, 3, b"def".as_slice() },
249        changed_declared_size = { (10, 1), 7, 3, b"def".as_slice() },
250    )]
251    fn rejected_retries_do_not_change_progress_or_latch_size(
252        id: SnapshotId,
253        size: i64,
254        position: i64,
255        chunk: &[u8],
256    ) {
257        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
258        start_three_byte_chunk(&mut s);
259        assert2::assert!(s.on_chunk(id, size, position, chunk) == SnapshotFetchStep::Restart);
260        assert2::assert!(s.next_position() == 3);
261    }
262
263    #[test]
264    fn accepted_chunks_complete_the_latched_snapshot() {
265        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
266        start_three_byte_chunk(&mut s);
267        assert2::assert!(
268            s.on_chunk((10, 1), 6, 3, b"def")
269                == SnapshotFetchStep::Complete(Bytes::from_static(b"abcdef"))
270        );
271    }
272
273    #[test]
274    fn fresh_transfer_can_accept_a_different_declared_size() {
275        let mut fresh = SnapshotFetchState::new((10, 1), NodeId(2));
276        assert2::assert!(fresh.on_chunk((10, 1), 3, 0, b"toolong") == SnapshotFetchStep::Restart);
277        assert2::assert!(
278            fresh.on_chunk((10, 1), 7, 0, b"toolong")
279                == SnapshotFetchStep::Complete(Bytes::from_static(b"toolong"))
280        );
281    }
282
283    #[test]
284    fn declared_size_over_cap_restarts() {
285        let mut s = SnapshotFetchState::new((10, 1), NodeId(2));
286        let too_big = METADATA_SNAPSHOT_FETCH_HARD_MAX.bytes_i64() + 1;
287        let step = s.on_chunk((10, 1), too_big, 0, b"abc");
288        assert2::assert!(step == SnapshotFetchStep::Restart);
289        assert2::assert!(s.next_position() == 0);
290    }
291
292    #[test]
293    fn reassembly_cap_is_exactly_one_gibibyte() {
294        // The cap bounds the memory a peer claiming to be the leader can force
295        // this follower to allocate, so its magnitude is the security property
296        // (finding M-3) — not just that some cap exists.
297        assert2::assert!(METADATA_SNAPSHOT_FETCH_HARD_MAX.bytes_i64() == 1_073_741_824);
298    }
299
300    #[test]
301    fn single_chunk_completes() {
302        let mut s = SnapshotFetchState::new((5, 0), NodeId(1));
303        match s.on_chunk((5, 0), 3, 0, b"xyz") {
304            SnapshotFetchStep::Complete(b) => assert2::assert!(b.as_ref() == b"xyz"),
305            other => panic!("expected Complete, got {other:?}"),
306        }
307    }
308}
krabka_kraft_core/snapshot_fetch.rs

krabka_kraft_core/
snapshot_fetch.rs
