use aion_core::{WorkflowListFilter, WorkflowListRequest, WorkflowSort};
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::Serialize;
use sha2::{Digest, Sha256};
use crate::StoreError;
#[derive(Serialize)]
struct QueryIdentity<'a> {
namespace: &'a str,
filter: &'a WorkflowListFilter,
sort: WorkflowSort,
}
const FINGERPRINT_LEN: usize = 8;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WorkflowCursor {
pub after_key: Vec<u8>,
}
impl WorkflowCursor {
pub fn mint(request: &WorkflowListRequest, last_page_key: &[u8]) -> Result<String, StoreError> {
let mut bytes = fingerprint(request)?.to_vec();
bytes.extend_from_slice(last_page_key);
Ok(URL_SAFE_NO_PAD.encode(bytes))
}
pub fn decode(request: &WorkflowListRequest) -> Result<Option<Self>, StoreError> {
let Some(encoded) = request.cursor.as_deref() else {
return Ok(None);
};
let bytes = URL_SAFE_NO_PAD.decode(encoded).map_err(|error| {
StoreError::InvalidQuery(format!("cursor is not a valid cursor token: {error}"))
})?;
if bytes.len() <= FINGERPRINT_LEN {
return Err(StoreError::InvalidQuery(String::from(
"cursor is truncated: it names no row",
)));
}
let (stamped, after_key) = bytes.split_at(FINGERPRINT_LEN);
if stamped != fingerprint(request)? {
return Err(StoreError::InvalidQuery(String::from(
"cursor was minted under a different namespace, filter, or sort; restart from the first page",
)));
}
Ok(Some(Self {
after_key: after_key.to_vec(),
}))
}
}
fn fingerprint(request: &WorkflowListRequest) -> Result<[u8; FINGERPRINT_LEN], StoreError> {
let identity = QueryIdentity {
namespace: &request.namespace,
filter: &request.filter,
sort: request.sort,
};
let canonical = serde_json::to_vec(&identity)
.map_err(|error| StoreError::Serialization(format!("cursor identity: {error}")))?;
let digest = Sha256::digest(&canonical);
let mut head = [0_u8; FINGERPRINT_LEN];
head.copy_from_slice(&digest[..FINGERPRINT_LEN]);
Ok(head)
}
#[cfg(test)]
mod tests {
use aion_core::{
SortDirection, WorkflowListFilter, WorkflowListRequest, WorkflowSort, WorkflowSortField,
WorkflowStatus,
};
use super::WorkflowCursor;
use crate::StoreError;
fn request(cursor: Option<String>) -> WorkflowListRequest {
WorkflowListRequest {
namespace: String::from("default"),
filter: WorkflowListFilter::default(),
sort: WorkflowSort {
field: WorkflowSortField::StartedAt,
direction: SortDirection::Desc,
},
cursor,
limit: 10,
}
}
#[test]
fn round_trips_under_the_same_query() -> Result<(), StoreError> {
let minted = WorkflowCursor::mint(&request(None), b"key-bytes")?;
let decoded = WorkflowCursor::decode(&request(Some(minted)))?;
assert_eq!(
decoded,
Some(WorkflowCursor {
after_key: b"key-bytes".to_vec()
})
);
Ok(())
}
#[test]
fn refuses_a_cursor_from_another_sort_or_filter() -> Result<(), StoreError> {
let minted = WorkflowCursor::mint(&request(None), b"k")?;
let mut other_sort = request(Some(minted.clone()));
other_sort.sort.direction = SortDirection::Asc;
assert!(matches!(
WorkflowCursor::decode(&other_sort),
Err(StoreError::InvalidQuery(_))
));
let mut other_filter = request(Some(minted));
other_filter.filter.statuses = vec![WorkflowStatus::Failed];
assert!(matches!(
WorkflowCursor::decode(&other_filter),
Err(StoreError::InvalidQuery(_))
));
Ok(())
}
#[test]
fn refuses_garbage_and_truncated_tokens() {
assert!(matches!(
WorkflowCursor::decode(&request(Some(String::from("***")))),
Err(StoreError::InvalidQuery(_))
));
assert!(matches!(
WorkflowCursor::decode(&request(Some(String::from("AAAA")))),
Err(StoreError::InvalidQuery(_))
));
}
#[test]
fn absent_cursor_is_the_first_page() -> Result<(), StoreError> {
assert_eq!(WorkflowCursor::decode(&request(None))?, None);
Ok(())
}
}