Skip to main content

fakecloud_core/
pagination.rs

1/// Offset-based pagination helper for AWS list operations.
2///
3/// Parses `next_token` as a numeric offset (defaulting to 0 if `None` or unparseable),
4/// slices `items` starting at that offset, and returns at most `max_results` items
5/// along with an optional next token for the following page.
6///
7/// Prefer [`paginate_checked`] for client-facing list ops: this variant treats a
8/// malformed `next_token` as offset 0, which can drive an infinite client
9/// pagination loop. It remains for callers whose service model declares no
10/// invalid-token error (returning one would be an undeclared error).
11#[must_use]
12pub fn paginate<T: Clone>(
13    items: &[T],
14    next_token: Option<&str>,
15    max_results: usize,
16) -> (Vec<T>, Option<String>) {
17    if max_results == 0 {
18        return (Vec::new(), None);
19    }
20    let offset: usize = next_token.and_then(|s| s.parse().ok()).unwrap_or(0);
21    let page = if offset < items.len() {
22        &items[offset..]
23    } else {
24        &[][..]
25    };
26    let has_more = page.len() > max_results;
27    let result: Vec<T> = page.iter().take(max_results).cloned().collect();
28    let token = if has_more {
29        Some((offset + max_results).to_string())
30    } else {
31        None
32    };
33    (result, token)
34}
35
36/// Error from [`paginate_checked`]: `next_token` was present but is not a valid
37/// offset token (not produced by a prior page of the same list op). AWS rejects
38/// such tokens with `InvalidNextToken` (or a service-specific equivalent);
39/// callers map this to their wire error (bug-audit 2026-05-28, 1.7).
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub struct InvalidNextToken;
42
43/// Strict variant of [`paginate`]: a `next_token` that is present but does not
44/// parse as a non-negative offset is rejected with [`InvalidNextToken`] instead
45/// of being silently treated as offset 0 (which can drive an infinite client
46/// pagination loop). `None` still means "first page".
47pub fn paginate_checked<T: Clone>(
48    items: &[T],
49    next_token: Option<&str>,
50    max_results: usize,
51) -> Result<(Vec<T>, Option<String>), InvalidNextToken> {
52    let offset: usize = match next_token {
53        None => 0,
54        Some(tok) => tok.parse().map_err(|_| InvalidNextToken)?,
55    };
56    if max_results == 0 {
57        return Ok((Vec::new(), None));
58    }
59    let page = if offset < items.len() {
60        &items[offset..]
61    } else {
62        &[][..]
63    };
64    let has_more = page.len() > max_results;
65    let result: Vec<T> = page.iter().take(max_results).cloned().collect();
66    let token = if has_more {
67        Some((offset + max_results).to_string())
68    } else {
69        None
70    };
71    Ok((result, token))
72}
73
74/// Parse a client-supplied offset token (as minted by [`paginate`] /
75/// [`paginate_checked`] / [`page_json_response`]). Absent or empty is offset 0.
76pub fn parse_offset_token(token: Option<&str>) -> Result<usize, InvalidNextToken> {
77    match token.filter(|t| !t.is_empty()) {
78        None => Ok(0),
79        Some(t) => t.parse().map_err(|_| InvalidNextToken),
80    }
81}
82
83/// Page a successful JSON list response in place: the array at `items` is
84/// sliced to `[start, start + size)` (`size` of `None` keeps everything from
85/// `start`), and `token_key` is set to the offset token of the next page when
86/// items remain, or removed when none do (AWS omits an exhausted token rather
87/// than sending it empty). Responses that are not JSON objects with an array
88/// at `items` are returned unchanged.
89///
90/// This lets a handler render its full, stably ordered listing and a service
91/// apply `MaxResults` / `NextToken` uniformly at its dispatch boundary.
92pub fn page_json_response(
93    resp: crate::service::AwsResponse,
94    items: &str,
95    token_key: &str,
96    start: usize,
97    size: Option<usize>,
98) -> crate::service::AwsResponse {
99    use crate::service::ResponseBody;
100    if !resp.status.is_success() {
101        return resp;
102    }
103    let ResponseBody::Bytes(bytes) = &resp.body else {
104        return resp;
105    };
106    let Ok(mut value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
107        return resp;
108    };
109    let Some(obj) = value.as_object_mut() else {
110        return resp;
111    };
112    let Some(list) = obj.get_mut(items).and_then(|v| v.as_array_mut()) else {
113        return resp;
114    };
115    let total = list.len();
116    let start = start.min(total);
117    let end = size.map_or(total, |n| start.saturating_add(n).min(total));
118    let page: Vec<serde_json::Value> = list.drain(start..end).collect();
119    *list = page;
120    if end < total {
121        obj.insert(
122            token_key.to_string(),
123            serde_json::Value::String(end.to_string()),
124        );
125    } else {
126        obj.remove(token_key);
127    }
128    crate::service::AwsResponse {
129        body: ResponseBody::Bytes(bytes::Bytes::from(
130            serde_json::to_vec(&value).expect("serde_json::Value serialization is infallible"),
131        )),
132        ..resp
133    }
134}
135
136/// Where a JSON-protocol paging input travels on the wire.
137#[derive(Clone, Copy, Debug, PartialEq, Eq)]
138pub enum PageLoc {
139    /// An `@httpQuery` parameter.
140    Query(&'static str),
141    /// A top-level JSON body member.
142    Body(&'static str),
143}
144
145/// One paginated JSON-protocol operation, as generated from its Smithy model
146/// by `scripts/generate-json-pagination-tables.py`.
147#[derive(Clone, Copy, Debug)]
148pub struct JsonPagedOp {
149    pub action: &'static str,
150    /// The input page token.
151    pub token: PageLoc,
152    /// The input page size, if the operation models one.
153    pub size: Option<PageLoc>,
154    /// The output list a page slices (its JSON member name).
155    pub items: &'static str,
156    /// The output member carrying the next page's token (JSON name).
157    pub next_token: &'static str,
158    /// The output token is `@required`: the last page sends it empty rather
159    /// than omitting it.
160    pub next_token_required: bool,
161    /// The error code the operation declares for a token it did not mint
162    /// (always an HTTP 400). `None` when it declares none, in which case an
163    /// unrecognised token starts from the top rather than returning an
164    /// undeclared error.
165    pub bad_token: Option<&'static str>,
166}
167
168/// A validated page request for a [`JsonPagedOp`].
169#[derive(Clone, Copy, Debug)]
170pub struct JsonPage {
171    op: &'static JsonPagedOp,
172    start: usize,
173    size: Option<usize>,
174}
175
176fn read_page_input(
177    req: &crate::service::AwsRequest,
178    body: &serde_json::Value,
179    loc: PageLoc,
180) -> Option<String> {
181    match loc {
182        PageLoc::Query(name) => req.query_params.get(name).cloned(),
183        PageLoc::Body(name) => match body.get(name)? {
184            serde_json::Value::String(s) => Some(s.clone()),
185            serde_json::Value::Number(n) => Some(n.to_string()),
186            _ => None,
187        },
188    }
189}
190
191/// Validate a request to one of `ops` before its handler runs: a token the
192/// service did not mint is rejected with the operation's declared error.
193/// Returns the window to slice the handler's listing to, or `None` when the
194/// operation is not paginated or the request asks for everything. A page size
195/// of 0 (the models' default) means "the service default", here the whole
196/// listing.
197pub fn validate_json_page(
198    ops: &'static [JsonPagedOp],
199    action: &str,
200    req: &crate::service::AwsRequest,
201) -> Result<Option<JsonPage>, crate::service::AwsServiceError> {
202    let Some(op) = ops.iter().find(|o| o.action == action) else {
203        return Ok(None);
204    };
205    let body: serde_json::Value = if req.body.is_empty() {
206        serde_json::Value::Null
207    } else {
208        serde_json::from_slice(&req.body).unwrap_or(serde_json::Value::Null)
209    };
210    let token = read_page_input(req, &body, op.token);
211    let start = match (parse_offset_token(token.as_deref()), op.bad_token) {
212        (Ok(n), _) => n,
213        (Err(_), Some(code)) => {
214            return Err(crate::service::AwsServiceError::aws_error(
215                http::StatusCode::BAD_REQUEST,
216                code,
217                format!("Invalid pagination token: {}", token.unwrap_or_default()),
218            ))
219        }
220        (Err(_), None) => 0,
221    };
222    let size = op
223        .size
224        .and_then(|loc| read_page_input(req, &body, loc))
225        .and_then(|v| v.parse::<usize>().ok())
226        .filter(|n| *n > 0);
227    if start == 0 && size.is_none() {
228        return Ok(None);
229    }
230    Ok(Some(JsonPage { op, start, size }))
231}
232
233/// Slice a handler's full listing to `page` (see [`page_json_response`]),
234/// keeping a `@required` output token present (empty) on the last page.
235pub fn apply_json_page(
236    resp: crate::service::AwsResponse,
237    page: JsonPage,
238) -> crate::service::AwsResponse {
239    let op = page.op;
240    let paged = page_json_response(resp, op.items, op.next_token, page.start, page.size);
241    if !op.next_token_required {
242        return paged;
243    }
244    let crate::service::ResponseBody::Bytes(bytes) = &paged.body else {
245        return paged;
246    };
247    let Ok(mut v) = serde_json::from_slice::<serde_json::Value>(bytes) else {
248        return paged;
249    };
250    match v.as_object_mut() {
251        Some(obj) if !obj.contains_key(op.next_token) => {
252            obj.insert(
253                op.next_token.to_string(),
254                serde_json::Value::String(String::new()),
255            );
256            crate::service::AwsResponse::json_value(paged.status, v)
257        }
258        _ => paged,
259    }
260}
261
262/// The operations in a Smithy model (`aws-models/<svc>.json`, parsed) whose
263/// input carries a page token (`NextToken` / `nextToken` / `position`) and
264/// whose output returns one, sorted. Services check their generated
265/// [`JsonPagedOp`] tables against this so a model refresh that adds a
266/// paginated operation fails a test until the table is regenerated.
267pub fn model_paginated_actions(model: &serde_json::Value) -> Vec<String> {
268    const TOKENS: [&str; 3] = ["NextToken", "nextToken", "position"];
269    let Some(shapes) = model["shapes"].as_object() else {
270        return Vec::new();
271    };
272    let members = |target: &serde_json::Value| {
273        target
274            .as_str()
275            .and_then(|t| shapes.get(t))
276            .map(|s| s["members"].clone())
277            .unwrap_or_default()
278    };
279    let mut out: Vec<String> = shapes
280        .iter()
281        .filter(|(_, s)| s["type"] == "operation")
282        .filter(|(_, s)| {
283            let input = members(&s["input"]["target"]);
284            let output = members(&s["output"]["target"]);
285            TOKENS.iter().any(|t| input.get(*t).is_some())
286                && TOKENS.iter().any(|t| output.get(*t).is_some())
287        })
288        .filter_map(|(id, _)| id.rsplit('#').next().map(str::to_string))
289        .collect();
290    out.sort_unstable();
291    out
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    fn json_body(resp: &crate::service::AwsResponse) -> serde_json::Value {
299        match &resp.body {
300            crate::service::ResponseBody::Bytes(b) => serde_json::from_slice(b).unwrap(),
301            _ => panic!("not bytes"),
302        }
303    }
304
305    #[test]
306    fn page_json_response_slices_and_sets_or_clears_the_token() {
307        let full = || {
308            crate::service::AwsResponse::ok_json(
309                serde_json::json!({"Things": [0, 1, 2, 3, 4], "NextToken": ""}),
310            )
311        };
312        let p1 = json_body(&page_json_response(
313            full(),
314            "Things",
315            "NextToken",
316            0,
317            Some(2),
318        ));
319        assert_eq!(p1["Things"], serde_json::json!([0, 1]));
320        assert_eq!(p1["NextToken"], "2");
321        let p3 = json_body(&page_json_response(
322            full(),
323            "Things",
324            "NextToken",
325            4,
326            Some(2),
327        ));
328        assert_eq!(p3["Things"], serde_json::json!([4]));
329        assert!(p3.get("NextToken").is_none(), "{p3}");
330        let all = json_body(&page_json_response(full(), "Things", "NextToken", 0, None));
331        assert_eq!(all["Things"].as_array().unwrap().len(), 5);
332        assert!(all.get("NextToken").is_none());
333        // A response without the list is left alone.
334        let other = json_body(&page_json_response(
335            crate::service::AwsResponse::ok_json(serde_json::json!({"X": 1})),
336            "Things",
337            "NextToken",
338            0,
339            Some(1),
340        ));
341        assert_eq!(other, serde_json::json!({"X": 1}));
342    }
343
344    fn req_with(query: &[(&str, &str)], body: serde_json::Value) -> crate::service::AwsRequest {
345        crate::service::AwsRequest {
346            service: "svc".into(),
347            action: "ListThings".into(),
348            region: "us-east-1".into(),
349            account_id: "000000000000".into(),
350            request_id: "rid".into(),
351            headers: http::HeaderMap::new(),
352            query_params: query
353                .iter()
354                .map(|(k, v)| (k.to_string(), v.to_string()))
355                .collect(),
356            body: bytes::Bytes::from(serde_json::to_vec(&body).unwrap()),
357            body_stream: parking_lot::Mutex::new(None),
358            path_segments: Vec::new(),
359            raw_path: "/".into(),
360            raw_query: String::new(),
361            method: http::Method::GET,
362            is_query_protocol: false,
363            access_key_id: None,
364            principal: None,
365        }
366    }
367
368    static OPS: [JsonPagedOp; 2] = [
369        JsonPagedOp {
370            action: "ListThings",
371            token: PageLoc::Query("nextToken"),
372            size: Some(PageLoc::Query("maxResults")),
373            items: "things",
374            next_token: "nextToken",
375            next_token_required: false,
376            bad_token: Some("BadRequestException"),
377        },
378        JsonPagedOp {
379            action: "DescribeThings",
380            token: PageLoc::Body("NextToken"),
381            size: Some(PageLoc::Body("MaxResults")),
382            items: "Things",
383            next_token: "NextToken",
384            next_token_required: true,
385            bad_token: None,
386        },
387    ];
388
389    #[test]
390    fn json_pages_validate_tokens_and_slice() {
391        let full =
392            |key: &str| crate::service::AwsResponse::ok_json(serde_json::json!({ key: [0, 1, 2] }));
393        // Query-carried token + size; a foreign token is the declared error.
394        let page = validate_json_page(
395            &OPS,
396            "ListThings",
397            &req_with(&[("maxResults", "2")], serde_json::json!({})),
398        )
399        .unwrap()
400        .unwrap();
401        let p1 = json_body(&apply_json_page(full("things"), page));
402        assert_eq!(p1["things"], serde_json::json!([0, 1]));
403        assert_eq!(p1["nextToken"], "2");
404        let err = validate_json_page(
405            &OPS,
406            "ListThings",
407            &req_with(&[("nextToken", "x")], serde_json::json!({})),
408        )
409        .unwrap_err();
410        assert_eq!(err.code(), "BadRequestException");
411        // Unpaginated op or "everything" requests are left alone.
412        assert!(
413            validate_json_page(&OPS, "Other", &req_with(&[], serde_json::json!({})))
414                .unwrap()
415                .is_none()
416        );
417        assert!(
418            validate_json_page(&OPS, "ListThings", &req_with(&[], serde_json::json!({})))
419                .unwrap()
420                .is_none()
421        );
422
423        // Body-carried; no declared token error means a foreign token starts
424        // over; a required token stays (empty) on the last page.
425        let page = validate_json_page(
426            &OPS,
427            "DescribeThings",
428            &req_with(
429                &[],
430                serde_json::json!({"MaxResults": 5, "NextToken": "junk"}),
431            ),
432        )
433        .unwrap()
434        .unwrap();
435        let last = json_body(&apply_json_page(full("Things"), page));
436        assert_eq!(last["Things"], serde_json::json!([0, 1, 2]));
437        assert_eq!(last["NextToken"], "");
438    }
439
440    #[test]
441    fn model_paginated_actions_needs_token_in_and_out() {
442        let model = serde_json::json!({"shapes": {
443            "s#A": {"type": "operation", "input": {"target": "s#AIn"}, "output": {"target": "s#AOut"}},
444            "s#AIn": {"type": "structure", "members": {"NextToken": {}}},
445            "s#AOut": {"type": "structure", "members": {"NextToken": {}, "Items": {}}},
446            "s#B": {"type": "operation", "input": {"target": "s#BIn"}, "output": {"target": "s#BOut"}},
447            "s#BIn": {"type": "structure", "members": {"nextToken": {}}},
448            "s#BOut": {"type": "structure", "members": {}}
449        }});
450        assert_eq!(model_paginated_actions(&model), vec!["A".to_string()]);
451    }
452
453    #[test]
454    fn parse_offset_token_rejects_non_offsets() {
455        assert_eq!(parse_offset_token(None), Ok(0));
456        assert_eq!(parse_offset_token(Some("")), Ok(0));
457        assert_eq!(parse_offset_token(Some("7")), Ok(7));
458        assert_eq!(parse_offset_token(Some("abc")), Err(InvalidNextToken));
459    }
460
461    #[test]
462    fn first_page() {
463        let items: Vec<i32> = (0..10).collect();
464        let (page, token) = paginate(&items, None, 3);
465        assert_eq!(page, vec![0, 1, 2]);
466        assert_eq!(token, Some("3".to_string()));
467    }
468
469    #[test]
470    fn middle_page() {
471        let items: Vec<i32> = (0..10).collect();
472        let (page, token) = paginate(&items, Some("3"), 3);
473        assert_eq!(page, vec![3, 4, 5]);
474        assert_eq!(token, Some("6".to_string()));
475    }
476
477    #[test]
478    fn last_page() {
479        let items: Vec<i32> = (0..10).collect();
480        let (page, token) = paginate(&items, Some("9"), 3);
481        assert_eq!(page, vec![9]);
482        assert_eq!(token, None);
483    }
484
485    #[test]
486    fn exact_page_boundary() {
487        let items: Vec<i32> = (0..6).collect();
488        let (page, token) = paginate(&items, Some("3"), 3);
489        assert_eq!(page, vec![3, 4, 5]);
490        assert_eq!(token, None);
491    }
492
493    #[test]
494    fn offset_beyond_items() {
495        let items: Vec<i32> = (0..3).collect();
496        let (page, token) = paginate(&items, Some("100"), 3);
497        assert!(page.is_empty());
498        assert_eq!(token, None);
499    }
500
501    #[test]
502    fn invalid_token_defaults_to_zero() {
503        let items: Vec<i32> = (0..5).collect();
504        let (page, token) = paginate(&items, Some("not_a_number"), 3);
505        assert_eq!(page, vec![0, 1, 2]);
506        assert_eq!(token, Some("3".to_string()));
507    }
508
509    #[test]
510    fn zero_max_results_returns_empty_page_without_token() {
511        // AWS list ops reject MaxResults=0 at the validation layer; if the helper
512        // ever sees zero it returns an empty page with no continuation token so
513        // callers can't accidentally paginate forever on a non-advancing offset.
514        let items: Vec<i32> = (0..5).collect();
515        let (page, token) = paginate(&items, None, 0);
516        assert!(page.is_empty());
517        assert_eq!(token, None);
518    }
519
520    #[test]
521    fn empty_items() {
522        let items: Vec<i32> = vec![];
523        let (page, token) = paginate(&items, None, 10);
524        assert!(page.is_empty());
525        assert_eq!(token, None);
526    }
527
528    // bug-audit 2026-05-28, 1.7: paginate_checked rejects a malformed next_token
529    // instead of silently treating it as offset 0.
530    #[test]
531    fn checked_none_is_first_page() {
532        let items: Vec<i32> = (0..5).collect();
533        let (page, token) = paginate_checked(&items, None, 3).unwrap();
534        assert_eq!(page, vec![0, 1, 2]);
535        assert_eq!(token, Some("3".to_string()));
536    }
537
538    #[test]
539    fn checked_valid_token_advances() {
540        let items: Vec<i32> = (0..5).collect();
541        let (page, token) = paginate_checked(&items, Some("3"), 3).unwrap();
542        assert_eq!(page, vec![3, 4]);
543        assert_eq!(token, None);
544    }
545
546    #[test]
547    fn checked_garbage_token_is_rejected() {
548        let items: Vec<i32> = (0..5).collect();
549        assert_eq!(
550            paginate_checked(&items, Some("not_a_number"), 3),
551            Err(InvalidNextToken)
552        );
553    }
554
555    #[test]
556    fn checked_negative_token_is_rejected() {
557        let items: Vec<i32> = (0..5).collect();
558        assert_eq!(
559            paginate_checked(&items, Some("-1"), 3),
560            Err(InvalidNextToken)
561        );
562    }
563}