From 2878de5b26a5980795107eb4dcd4f41ac162e255 Mon Sep 17 00:00:00 2001
From: "Hargobind S. Khalsa" <khalsah@gmail.com>
Date: Sat, 31 Oct 2020 13:40:50 -0700
Subject: [PATCH] Validate the query hash before persisting queries

---
 Cargo.toml                                 |  3 ++-
 src/extensions/apollo_persisted_queries.rs | 17 +++++++++++------
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9a83a37a..e68d2094 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,7 +28,7 @@ default = [
 	"uuid",
 ]
 apollo_tracing = ["chrono"]
-apollo_persisted_queries = ["async-mutex", "lru"]
+apollo_persisted_queries = ["async-mutex", "lru", "sha2"]
 multipart = ["bytes", "multer", "tempfile"]
 unblock = ["blocking"]
 string_number = ["num-traits"]
@@ -70,6 +70,7 @@ bytes = { version = "0.5.4", optional = true }
 lru = { version = "0.6.0", optional = true }
 multer = { version = "1.2.2", optional = true }
 num-traits = { version = "0.2.12", optional = true }
+sha2 = { version = "0.9.1", optional = true }
 tempfile = { version = "3.1.0", optional = true }
 
 [dev-dependencies]
diff --git a/src/extensions/apollo_persisted_queries.rs b/src/extensions/apollo_persisted_queries.rs
index 6148de33..ffc5bc9d 100644
--- a/src/extensions/apollo_persisted_queries.rs
+++ b/src/extensions/apollo_persisted_queries.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 use async_mutex::Mutex;
 use serde::Deserialize;
+use sha2::{Digest, Sha256};
 
 use crate::extensions::{Extension, ExtensionContext, ExtensionFactory};
 use crate::{from_value, Request, ServerError, ServerResult};
@@ -98,10 +99,14 @@ impl<T: CacheStorage> Extension for ApolloPersistedQueriesExtension<T> {
                     Err(ServerError::new("PersistedQueryNotFound".to_string()))
                 }
             } else {
-                self.storage
-                    .set(persisted_query.sha256_hash, request.query.clone())
-                    .await;
-                Ok(request)
+                let sha256_hash = format!("{:x}", Sha256::digest(request.query.as_bytes()));
+
+                if (persisted_query.sha256_hash != sha256_hash) {
+                    Err(ServerError::new("provided sha does not match query"))
+                } else {
+                    self.storage.set(sha256_hash, request.query.clone()).await;
+                    Ok(request)
+                }
             }
         } else {
             Ok(request)
@@ -134,7 +139,7 @@ mod tests {
             "persistedQuery".to_string(),
             value!({
                 "version": 1,
-                "sha256Hash": "abc",
+                "sha256Hash": "854174ebed716fe24fd6659c30290aecd9bc1d17dc4f47939a1848a1b8ed3c6b",
             }),
         );
 
@@ -150,7 +155,7 @@ mod tests {
             "persistedQuery".to_string(),
             value!({
                 "version": 1,
-                "sha256Hash": "abc",
+                "sha256Hash": "854174ebed716fe24fd6659c30290aecd9bc1d17dc4f47939a1848a1b8ed3c6b",
             }),
         );