Skip to main content

autumn_plugin_graphql/
persisted.rs

1//! Persisted operations: Automatic Persisted Queries (APQ) and trusted
2//! documents.
3//!
4//! - **APQ** ([`PersistedQueryMode::Automatic`](crate::config::PersistedQueryMode)):
5//!   clients send `extensions.persistedQuery.sha256Hash` instead of the
6//!   document. On a miss the server answers `PERSISTED_QUERY_NOT_FOUND`, the
7//!   client re-sends hash + document, and the server verifies the hash before
8//!   caching it in a [`PersistedQueryStore`]. Saves bandwidth; not a security
9//!   boundary.
10//! - **Trusted documents**
11//!   ([`PersistedQueryMode::Trusted`](crate::config::PersistedQueryMode)): only
12//!   documents in a manifest produced at client build time may run. Anything
13//!   else is refused with `OPERATION_NOT_ALLOWLISTED`. This *is* a security
14//!   boundary: an attacker cannot run a query your clients never shipped.
15//!
16//! Manifests may be Apollo's `persisted-query-manifest.json`
17//! (`{"format":"apollo-persisted-query-manifest","version":1,"operations":[{"id","body",…}]}`)
18//! or a flat `{"<sha256 hex>": "<document>"}` map. Requests may name a
19//! document by the APQ extension or by a top-level `documentId` (the
20//! GraphQL-over-HTTP persisted-documents proposal; a `sha256:` prefix is
21//! accepted).
22
23use 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
33/// A boxed, `Send` future — the shape store methods return so the trait
34/// stays object-safe.
35pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
36
37/// Where APQ registrations live.
38///
39/// The in-memory [`InMemoryPersistedQueryStore`] is per process; a fleet
40/// behind a load balancer shares one by implementing this trait over Redis
41/// or the database (a miss is harmless — the client re-sends the document —
42/// so a best-effort store is fine).
43pub trait PersistedQueryStore: Send + Sync + 'static {
44    /// The document registered under `hash`, if any.
45    fn get<'a>(&'a self, hash: &'a str) -> BoxFuture<'a, Option<Arc<str>>>;
46    /// Register `document` under `hash`. The hash has already been verified.
47    fn put<'a>(&'a self, hash: &'a str, document: Arc<str>) -> BoxFuture<'a, ()>;
48}
49
50/// A bounded, least-recently-used in-memory [`PersistedQueryStore`].
51pub struct InMemoryPersistedQueryStore {
52    cache: Mutex<lru::LruCache<String, Arc<str>>>,
53}
54
55impl InMemoryPersistedQueryStore {
56    /// A store holding at most `capacity` documents (at least one).
57    #[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    /// Number of cached documents.
66    #[must_use]
67    pub fn len(&self) -> usize {
68        self.cache.lock().map_or(0, |cache| cache.len())
69    }
70
71    /// Whether the cache is empty.
72    #[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        // A poisoned lock only means another request panicked mid-update;
89        // the LRU is still structurally sound, so keep serving from it.
90        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/// Lower-case hex SHA-256 of `text` — the APQ document hash.
109#[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/// Whether `id` looks like a SHA-256 hex digest.
121#[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/// A manifest could not be loaded.
127#[derive(Debug, thiserror::Error)]
128pub enum ManifestError {
129    /// The file could not be read.
130    #[error("cannot read persisted-query manifest {path}: {source}")]
131    Io {
132        /// The manifest path.
133        path: String,
134        /// The underlying error.
135        source: std::io::Error,
136    },
137    /// The file is not a manifest in a supported format.
138    #[error("invalid persisted-query manifest: {0}")]
139    Invalid(String),
140}
141
142/// The allow-list of trusted documents.
143///
144/// Documents are addressable two ways: by the manifest's operation `id`
145/// (what a client sends), and by the SHA-256 of their text (so a client that
146/// sends the full document of a trusted operation is accepted too).
147#[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    /// An empty allow-list.
172    #[must_use]
173    pub fn new() -> Self {
174        Self::default()
175    }
176
177    /// Add a document, addressed by its SHA-256.
178    #[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    /// Add a document under an explicit manifest `id`.
187    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    /// Parse a manifest (Apollo format or a flat `{id: document}` map).
196    ///
197    /// # Errors
198    ///
199    /// Returns [`ManifestError::Invalid`] for malformed JSON, an unknown
200    /// Apollo manifest format/version, or a flat entry whose SHA-256 id does
201    /// not match its document.
202    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    /// Load a manifest file.
237    ///
238    /// # Errors
239    ///
240    /// See [`from_json`](Self::from_json); plus [`ManifestError::Io`].
241    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    /// Merge another allow-list into this one.
250    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    /// The document with manifest id (or SHA-256) `id`.
256    #[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    /// Whether a document with this exact text is trusted.
265    #[must_use]
266    pub fn contains_document(&self, document: &str) -> bool {
267        self.by_hash.contains_key(&sha256_hex(document))
268    }
269
270    /// Number of distinct trusted documents.
271    #[must_use]
272    pub fn len(&self) -> usize {
273        self.by_hash.len()
274    }
275
276    /// Whether the allow-list is empty.
277    #[must_use]
278    pub fn is_empty(&self) -> bool {
279        self.by_hash.is_empty()
280    }
281}
282
283/// Strip the `sha256:` prefix the persisted-documents proposal allows.
284fn 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}