1#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub struct InvalidNextToken;
42
43pub 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
74pub 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
83pub 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
138pub enum PageLoc {
139 Query(&'static str),
141 Body(&'static str),
143}
144
145#[derive(Clone, Copy, Debug)]
148pub struct JsonPagedOp {
149 pub action: &'static str,
150 pub token: PageLoc,
152 pub size: Option<PageLoc>,
154 pub items: &'static str,
156 pub next_token: &'static str,
158 pub next_token_required: bool,
161 pub bad_token: Option<&'static str>,
166}
167
168#[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
191pub 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
233pub 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
262pub 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 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 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 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 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 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 #[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}