1use std::collections::BTreeSet;
2use std::error::Error;
3use std::fmt::{self, Display, Formatter};
4
5use lodestone_wire::{
6 Cursor, ExportUsageRequest, ExportUsageResult, ListIndexesRequest, ListIndexesResult, SearchRequest, SearchResult,
7};
8
9use crate::config::parse_u64;
10use crate::{CallOptions, Client, ClientError};
11
12const MAX_PAGES: u32 = 10_000;
13
14pub struct IndexPager {
15 client: Client,
16 request: ListIndexesRequest,
17 options: CallOptions,
18 seen_cursors: BTreeSet<String>,
19 pages: u32,
20 ended: bool,
21}
22
23pub struct SearchPager {
24 client: Client,
25 request: SearchRequest,
26 options: CallOptions,
27 seen_cursors: BTreeSet<String>,
28 pages: u32,
29 ended: bool,
30}
31
32pub struct UsagePager {
33 client: Client,
34 request: ExportUsageRequest,
35 options: CallOptions,
36 seen_cursors: BTreeSet<String>,
37 prior_cumulative_through: Option<u64>,
38 prior_sequence: Option<u64>,
39 prior_authority_revision: Option<u64>,
40 pages: u32,
41 ended: bool,
42}
43
44impl Client {
45 pub fn index_pager(&self, request: ListIndexesRequest, options: CallOptions) -> Result<IndexPager, PagerError> {
46 validate_page_size(request.page_size)?;
47 Ok(IndexPager {
48 client: self.clone(),
49 seen_cursors: cursor_history(&request.cursor)?,
50 request,
51 options,
52 pages: 0,
53 ended: false,
54 })
55 }
56
57 pub fn search_pager(&self, request: SearchRequest, options: CallOptions) -> Result<SearchPager, PagerError> {
58 validate_page_size(request.page_size)?;
59 let valid_binding = match (&request.indexes, &request.pit) {
60 (Some(indexes), None) => !indexes.is_empty(),
61 (None, Some(_)) => true,
62 _ => false,
63 };
64 if !valid_binding {
65 return Err(PagerError::InvalidBinding);
66 }
67 Ok(SearchPager {
68 client: self.clone(),
69 seen_cursors: cursor_history(&request.cursor)?,
70 request,
71 options,
72 pages: 0,
73 ended: false,
74 })
75 }
76
77 pub fn usage_pager(&self, request: ExportUsageRequest, options: CallOptions) -> Result<UsagePager, PagerError> {
78 validate_page_size(request.page_size)?;
79 let start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
80 let end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
81 if start >= end {
82 return Err(PagerError::InvalidBinding);
83 }
84 Ok(UsagePager {
85 client: self.clone(),
86 seen_cursors: cursor_history(&request.cursor)?,
87 request,
88 options,
89 prior_cumulative_through: None,
90 prior_sequence: None,
91 prior_authority_revision: None,
92 pages: 0,
93 ended: false,
94 })
95 }
96}
97
98impl IndexPager {
99 pub async fn next_page(&mut self) -> Result<Option<ListIndexesResult>, PagerError> {
100 if self.ended {
101 return Ok(None);
102 }
103 consume_budget(&mut self.pages)?;
104 let result = self
105 .client
106 .list_indexes(self.request.clone(), self.options.clone())
107 .await
108 .map_err(PagerError::Client)?;
109 if result.restart_count != self.request.restart_count {
110 self.ended = true;
111 return Err(PagerError::StaleCursor);
112 }
113 advance_cursor(
114 &mut self.request.cursor,
115 &mut self.seen_cursors,
116 result.next_cursor.clone(),
117 &mut self.ended,
118 )?;
119 Ok(Some(result))
120 }
121}
122
123impl SearchPager {
124 pub async fn next_page(&mut self) -> Result<Option<SearchResult>, PagerError> {
125 if self.ended {
126 return Ok(None);
127 }
128 consume_budget(&mut self.pages)?;
129 let result = self
130 .client
131 .search(self.request.clone(), self.options.clone())
132 .await
133 .map_err(PagerError::Client)?;
134 advance_cursor(
135 &mut self.request.cursor,
136 &mut self.seen_cursors,
137 result.next_cursor.clone(),
138 &mut self.ended,
139 )?;
140 Ok(Some(result))
141 }
142}
143
144impl UsagePager {
145 pub async fn next_page(&mut self) -> Result<Option<ExportUsageResult>, PagerError> {
146 if self.ended {
147 return Ok(None);
148 }
149 consume_budget(&mut self.pages)?;
150 let result = self
151 .client
152 .export_usage(self.request.clone(), self.options.clone())
153 .await
154 .map_err(PagerError::Client)?;
155 if let Err(error) = validate_usage_page(
156 &self.request,
157 &result,
158 &mut self.prior_cumulative_through,
159 &mut self.prior_sequence,
160 &mut self.prior_authority_revision,
161 ) {
162 self.ended = true;
163 return Err(error);
164 }
165 advance_cursor(
166 &mut self.request.cursor,
167 &mut self.seen_cursors,
168 result.next_cursor.clone(),
169 &mut self.ended,
170 )?;
171 Ok(Some(result))
172 }
173}
174
175fn validate_usage_page(
176 request: &ExportUsageRequest,
177 result: &ExportUsageResult,
178 prior_cumulative_through: &mut Option<u64>,
179 prior_sequence: &mut Option<u64>,
180 prior_authority_revision: &mut Option<u64>,
181) -> Result<(), PagerError> {
182 let request_start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
183 let request_end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
184 let cumulative = parse_u64(&result.cumulative_through.value).ok_or(PagerError::InvalidCorePage)?;
185 if prior_cumulative_through.is_some_and(|prior| cumulative < prior) {
186 return Err(PagerError::AuthorityRegression);
187 }
188
189 let mut rows = Vec::with_capacity(result.units.len().saturating_add(result.corrections.len()));
190 for row in &result.units {
191 rows.push(validate_usage_row(
192 &row.sequence.value,
193 &row.window_start.value,
194 &row.window_end.value,
195 &row.authority_revision.value,
196 &row.source_digest,
197 request_start,
198 request_end,
199 )?);
200 }
201 for row in &result.corrections {
202 if parse_u64(&row.correction_sequence.value).is_none_or(|value| value == 0) {
203 return Err(PagerError::InvalidCorePage);
204 }
205 rows.push(validate_usage_row(
206 &row.sequence.value,
207 &row.window_start.value,
208 &row.window_end.value,
209 &row.authority_revision.value,
210 &row.source_digest,
211 request_start,
212 request_end,
213 )?);
214 }
215 rows.sort_unstable_by_key(|row| row.0);
216 for (sequence, authority_revision) in rows {
217 if prior_sequence.is_some_and(|prior| sequence <= prior)
218 || prior_authority_revision.is_some_and(|prior| authority_revision != prior)
219 {
220 return Err(PagerError::AuthorityRegression);
221 }
222 *prior_sequence = Some(sequence);
223 *prior_authority_revision = Some(authority_revision);
224 }
225 if prior_sequence.is_some_and(|sequence| cumulative < sequence)
226 || result.next_cursor.is_some() && result.units.is_empty() && result.corrections.is_empty()
227 {
228 return Err(PagerError::InvalidCorePage);
229 }
230 *prior_cumulative_through = Some(cumulative);
231 Ok(())
232}
233
234fn validate_usage_row(
235 sequence: &str,
236 window_start: &str,
237 window_end: &str,
238 authority_revision: &str,
239 source_digest: &lodestone_wire::Digest,
240 request_start: u64,
241 request_end: u64,
242) -> Result<(u64, u64), PagerError> {
243 let sequence = parse_u64(sequence).ok_or(PagerError::InvalidCorePage)?;
244 let window_start = parse_u64(window_start).ok_or(PagerError::InvalidCorePage)?;
245 let window_end = parse_u64(window_end).ok_or(PagerError::InvalidCorePage)?;
246 let authority_revision = parse_u64(authority_revision).ok_or(PagerError::InvalidCorePage)?;
247 if sequence == 0
248 || authority_revision == 0
249 || window_start < request_start
250 || window_end > request_end
251 || window_start >= window_end
252 || source_digest.validate().is_err()
253 {
254 return Err(PagerError::InvalidCorePage);
255 }
256 Ok((sequence, authority_revision))
257}
258
259fn validate_page_size(page_size: u32) -> Result<(), PagerError> {
260 if page_size == 0 || page_size > 10_000 {
261 Err(PagerError::InvalidBinding)
262 } else {
263 Ok(())
264 }
265}
266
267fn consume_budget(pages: &mut u32) -> Result<(), PagerError> {
268 if *pages >= MAX_PAGES {
269 return Err(PagerError::PageBudgetExhausted);
270 }
271 *pages = pages.saturating_add(1);
272 Ok(())
273}
274
275fn advance_cursor(
276 request_cursor: &mut Option<Cursor>,
277 seen_cursors: &mut BTreeSet<String>,
278 next_cursor: Option<Cursor>,
279 ended: &mut bool,
280) -> Result<(), PagerError> {
281 if let Some(next) = &next_cursor {
282 if !valid_cursor(next) || !seen_cursors.insert(next.value.clone()) {
283 *ended = true;
284 return Err(PagerError::CursorCycle);
285 }
286 } else {
287 *ended = true;
288 }
289 *request_cursor = next_cursor;
290 Ok(())
291}
292
293fn cursor_history(cursor: &Option<Cursor>) -> Result<BTreeSet<String>, PagerError> {
294 let mut seen = BTreeSet::new();
295 if let Some(cursor) = cursor {
296 if !valid_cursor(cursor) {
297 return Err(PagerError::InvalidBinding);
298 }
299 let _ = seen.insert(cursor.value.clone());
300 }
301 Ok(seen)
302}
303
304fn valid_cursor(cursor: &Cursor) -> bool {
305 !cursor.value.is_empty()
306 && cursor.value.len() <= 16 * 1024
307 && !cursor.value.bytes().any(|byte| byte.is_ascii_control())
308}
309
310#[derive(Debug)]
311pub enum PagerError {
312 InvalidBinding,
313 Client(ClientError),
314 StaleCursor,
315 CursorCycle,
316 AuthorityRegression,
317 InvalidCorePage,
318 PageBudgetExhausted,
319}
320
321impl Display for PagerError {
322 fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
323 formatter.write_str(match self {
324 Self::InvalidBinding => "pager filter or page binding is invalid",
325 Self::Client(_) => "Core page request failed",
326 Self::StaleCursor => "Core catalog cursor became stale",
327 Self::CursorCycle => "Core pagination cursor repeated or is invalid",
328 Self::AuthorityRegression => "Core paged authority watermark regressed",
329 Self::InvalidCorePage => "Core page violates its generated semantic contract",
330 Self::PageBudgetExhausted => "Core pager reached its bounded page budget",
331 })
332 }
333}
334
335impl Error for PagerError {}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340
341 fn cursor(value: &str) -> Cursor {
342 Cursor {
343 value: value.to_owned(),
344 }
345 }
346
347 #[test]
348 fn cursor_history_refuses_non_adjacent_cycles() -> Result<(), PagerError> {
349 let mut request = None;
350 let mut seen = BTreeSet::new();
351 let mut ended = false;
352 advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended)?;
353 advance_cursor(&mut request, &mut seen, Some(cursor("b")), &mut ended)?;
354 assert!(matches!(
355 advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended),
356 Err(PagerError::CursorCycle)
357 ));
358 assert!(ended);
359 Ok(())
360 }
361
362 #[test]
363 fn initial_cursor_is_validated_and_participates_in_cycle_detection() -> Result<(), PagerError> {
364 assert!(matches!(
365 cursor_history(&Some(cursor("bad\nvalue"))),
366 Err(PagerError::InvalidBinding)
367 ));
368 let mut request = Some(cursor("initial"));
369 let mut seen = cursor_history(&request)?;
370 let mut ended = false;
371 assert!(matches!(
372 advance_cursor(&mut request, &mut seen, Some(cursor("initial")), &mut ended),
373 Err(PagerError::CursorCycle)
374 ));
375 Ok(())
376 }
377
378 #[test]
379 fn usage_sequence_watermark_is_not_compared_to_time_window() -> Result<(), PagerError> {
380 let request = ExportUsageRequest {
381 window_start: lodestone_wire::U64 { value: "10".to_owned() },
382 window_end: lodestone_wire::U64 { value: "20".to_owned() },
383 page_size: 10,
384 cursor: None,
385 };
386 let result = ExportUsageResult {
387 units: Vec::new(),
388 corrections: Vec::new(),
389 next_cursor: None,
390 cumulative_through: lodestone_wire::U64 {
391 value: "1000000".to_owned(),
392 },
393 };
394 validate_usage_page(&request, &result, &mut None, &mut None, &mut None)?;
395 Ok(())
396 }
397}