autumn_plugin_graphql/
persisted.rs1use std::collections::HashMap;
24use std::future::Future;
25use std::num::NonZeroUsize;
26use std::path::Path;
27use std::pin::Pin;
28use std::sync::{Arc, Mutex};
29
30use serde::Deserialize;
31use sha2::{Digest, Sha256};
32
33pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
36
37pub trait PersistedQueryStore: Send + Sync + 'static {
44 fn get<'a>(&'a self, hash: &'a str) -> BoxFuture<'a, Option<Arc<str>>>;
46 fn put<'a>(&'a self, hash: &'a str, document: Arc<str>) -> BoxFuture<'a, ()>;
48}
49
50pub struct InMemoryPersistedQueryStore {
52 cache: Mutex<lru::LruCache<String, Arc<str>>>,
53}
54
55impl InMemoryPersistedQueryStore {
56 #[must_use]
58 pub fn new(capacity: usize) -> Self {
59 let capacity = NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::MIN);
60 Self {
61 cache: Mutex::new(lru::LruCache::new(capacity)),
62 }
63 }
64
65 #[must_use]
67 pub fn len(&self) -> usize {
68 self.cache.lock().map_or(0, |cache| cache.len())
69 }
70
71 #[must_use]
73 pub fn is_empty(&self) -> bool {
74 self.len() == 0
75 }
76}
77
78impl std::fmt::Debug for InMemoryPersistedQueryStore {
79 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80 f.debug_struct("InMemoryPersistedQueryStore")
81 .field("len", &self.len())
82 .finish()
83 }
84}
85
86impl PersistedQueryStore for InMemoryPersistedQueryStore {
87 fn get<'a>(&'a self, hash: &'a str) -> BoxFuture<'a, Option<Arc<str>>> {
88 let found = self
91 .cache
92 .lock()
93 .unwrap_or_else(std::sync::PoisonError::into_inner)
94 .get(hash)
95 .cloned();
96 Box::pin(std::future::ready(found))
97 }
98
99 fn put<'a>(&'a self, hash: &'a str, document: Arc<str>) -> BoxFuture<'a, ()> {
100 self.cache
101 .lock()
102 .unwrap_or_else(std::sync::PoisonError::into_inner)
103 .put(hash.to_owned(), document);
104 Box::pin(std::future::ready(()))
105 }
106}
107
108#[must_use]
110pub fn sha256_hex(text: &str) -> String {
111 use std::fmt::Write as _;
112 let digest = Sha256::digest(text.as_bytes());
113 let mut out = String::with_capacity(64);
114 for byte in digest {
115 let _ = write!(out, "{byte:02x}");
116 }
117 out
118}
119
120#[must_use]
122pub fn is_sha256_hex(id: &str) -> bool {
123 id.len() == 64 && id.bytes().all(|b| b.is_ascii_hexdigit())
124}
125
126#[derive(Debug, thiserror::Error)]
128pub enum ManifestError {
129 #[error("cannot read persisted-query manifest {path}: {source}")]
131 Io {
132 path: String,
134 source: std::io::Error,
136 },
137 #[error("invalid persisted-query manifest: {0}")]
139 Invalid(String),
140}
141
142#[derive(Debug, Clone, Default)]
148pub struct TrustedDocuments {
149 by_id: HashMap<String, Arc<str>>,
150 by_hash: HashMap<String, Arc<str>>,
151}
152
153#[derive(Deserialize)]
154#[serde(untagged)]
155enum ManifestFile {
156 Apollo {
157 format: String,
158 version: u32,
159 operations: Vec<ManifestOperation>,
160 },
161 Flat(HashMap<String, String>),
162}
163
164#[derive(Deserialize)]
165struct ManifestOperation {
166 id: String,
167 body: String,
168}
169
170impl TrustedDocuments {
171 #[must_use]
173 pub fn new() -> Self {
174 Self::default()
175 }
176
177 #[must_use]
179 pub fn with_document(mut self, document: impl Into<String>) -> Self {
180 let document: String = document.into();
181 let hash = sha256_hex(&document);
182 self.insert(hash, document);
183 self
184 }
185
186 pub fn insert(&mut self, id: impl Into<String>, document: impl Into<String>) {
188 let document: Arc<str> = Arc::from(document.into());
189 self.by_hash
190 .insert(sha256_hex(&document), Arc::clone(&document));
191 self.by_id
192 .insert(normalize_id(&id.into()).to_owned(), document);
193 }
194
195 pub fn from_json(text: &str) -> Result<Self, ManifestError> {
203 let parsed: ManifestFile =
204 serde_json::from_str(text).map_err(|e| ManifestError::Invalid(e.to_string()))?;
205 let mut documents = Self::new();
206 match parsed {
207 ManifestFile::Apollo {
208 format,
209 version,
210 operations,
211 } => {
212 if format != "apollo-persisted-query-manifest" || version != 1 {
213 return Err(ManifestError::Invalid(format!(
214 "unsupported manifest format {format:?} version {version}"
215 )));
216 }
217 for op in operations {
218 documents.insert(op.id, op.body);
219 }
220 }
221 ManifestFile::Flat(map) => {
222 for (id, body) in map {
223 let id = normalize_id(&id).to_ascii_lowercase();
224 if is_sha256_hex(&id) && sha256_hex(&body) != id {
225 return Err(ManifestError::Invalid(format!(
226 "entry {id} does not hash to its id"
227 )));
228 }
229 documents.insert(id, body);
230 }
231 }
232 }
233 Ok(documents)
234 }
235
236 pub fn from_file(path: &Path) -> Result<Self, ManifestError> {
242 let text = std::fs::read_to_string(path).map_err(|source| ManifestError::Io {
243 path: path.display().to_string(),
244 source,
245 })?;
246 Self::from_json(&text)
247 }
248
249 pub fn extend(&mut self, other: Self) {
251 self.by_id.extend(other.by_id);
252 self.by_hash.extend(other.by_hash);
253 }
254
255 #[must_use]
257 pub fn get(&self, id: &str) -> Option<&Arc<str>> {
258 let id = normalize_id(id);
259 self.by_id
260 .get(id)
261 .or_else(|| self.by_hash.get(&id.to_ascii_lowercase()))
262 }
263
264 #[must_use]
266 pub fn contains_document(&self, document: &str) -> bool {
267 self.by_hash.contains_key(&sha256_hex(document))
268 }
269
270 #[must_use]
272 pub fn len(&self) -> usize {
273 self.by_hash.len()
274 }
275
276 #[must_use]
278 pub fn is_empty(&self) -> bool {
279 self.by_hash.is_empty()
280 }
281}
282
283fn normalize_id(id: &str) -> &str {
285 id.strip_prefix("sha256:").unwrap_or(id)
286}
287
288#[cfg(test)]
289#[allow(clippy::unwrap_used)]
290mod tests {
291 use super::*;
292
293 const DOC: &str = "{ notes { id } }";
294
295 #[test]
296 fn sha256_matches_known_vector() {
297 assert_eq!(
298 sha256_hex("abc"),
299 "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
300 );
301 assert!(is_sha256_hex(&sha256_hex(DOC)));
302 assert!(!is_sha256_hex("xyz"));
303 }
304
305 #[tokio::test]
306 async fn lru_store_evicts_least_recently_used() {
307 let store = InMemoryPersistedQueryStore::new(2);
308 store.put("a", Arc::from("A")).await;
309 store.put("b", Arc::from("B")).await;
310 assert_eq!(store.get("a").await.as_deref(), Some("A"));
311 store.put("c", Arc::from("C")).await;
312 assert_eq!(store.get("b").await, None, "b was least recently used");
313 assert_eq!(store.len(), 2);
314 assert!(!store.is_empty());
315 assert!(InMemoryPersistedQueryStore::new(0).is_empty());
316 }
317
318 #[test]
319 fn apollo_manifest_is_addressable_by_id_and_by_hash() {
320 let manifest = format!(
321 r#"{{"format":"apollo-persisted-query-manifest","version":1,
322 "operations":[{{"id":"notes-v1","name":"Notes","type":"query","body":"{DOC}"}}]}}"#
323 );
324 let docs = TrustedDocuments::from_json(&manifest).unwrap();
325 assert_eq!(docs.get("notes-v1").map(|d| &**d), Some(DOC));
326 assert_eq!(docs.get(&sha256_hex(DOC)).map(|d| &**d), Some(DOC));
327 assert!(docs.contains_document(DOC));
328 assert!(!docs.contains_document("{ other }"));
329 assert_eq!(docs.len(), 1);
330 }
331
332 #[test]
333 fn flat_manifest_is_verified() {
334 let hash = sha256_hex(DOC);
335 let docs = TrustedDocuments::from_json(&format!(r#"{{"sha256:{hash}":"{DOC}"}}"#)).unwrap();
336 assert!(docs.get(&format!("sha256:{hash}")).is_some());
337 assert!(docs.get(&hash.to_ascii_uppercase()).is_some());
338
339 let wrong = sha256_hex("{ other }");
340 assert!(TrustedDocuments::from_json(&format!(r#"{{"{wrong}":"{DOC}"}}"#)).is_err());
341 }
342
343 #[test]
344 fn bad_manifests_are_rejected() {
345 assert!(TrustedDocuments::from_json("[]").is_err());
346 assert!(
347 TrustedDocuments::from_json(
348 r#"{"format":"apollo-persisted-query-manifest","version":2,"operations":[]}"#
349 )
350 .is_err()
351 );
352 assert!(TrustedDocuments::from_file(Path::new("/nonexistent/manifest.json")).is_err());
353 }
354
355 #[test]
356 fn documents_can_be_added_in_code_and_merged() {
357 let mut docs = TrustedDocuments::new().with_document(DOC);
358 let mut other = TrustedDocuments::new();
359 other.insert("custom", "{ x }");
360 docs.extend(other);
361 assert_eq!(docs.len(), 2);
362 assert!(docs.get("custom").is_some());
363 assert!(!docs.is_empty());
364 }
365}