fix: handle subscriptions domains & special query keys

graph-gateway/src/auth.rs

@@ -1,11 +1,7 @@
+use std::collections::{HashMap, HashSet};
+use std::sync::atomic::{self, AtomicUsize};
+use std::sync::Arc;
 use std::time::Duration;
-use std::{
-    collections::{HashMap, HashSet},
-    sync::{
-        atomic::{self, AtomicUsize},
-        Arc,
-    },
-};
 
 use alloy_primitives::Address;
 use anyhow::{anyhow, bail, ensure, Result};
@@ -22,8 +18,10 @@ use crate::topology::Deployment;
 pub struct AuthHandler {
     pub api_keys: Eventual<Ptr<HashMap<String, Arc<APIKey>>>>,
     pub special_api_keys: HashSet<String>,
+    pub special_query_key_signers: HashSet<Address>,
     pub api_key_payment_required: bool,
     pub subscriptions: Eventual<Ptr<HashMap<Address, Subscription>>>,
+    pub subscription_domains: HashMap<u64, Address>,
     pub subscription_query_counters: RwLock<HashMap<Address, AtomicUsize>>,
 }
 
@@ -43,14 +41,18 @@ impl AuthHandler {
     pub fn create(
         api_keys: Eventual<Ptr<HashMap<String, Arc<APIKey>>>>,
         special_api_keys: HashSet<String>,
+        special_query_key_signers: HashSet<Address>,
         api_key_payment_required: bool,
         subscriptions: Eventual<Ptr<HashMap<Address, Subscription>>>,
+        subscription_domains: HashMap<u64, Address>,
     ) -> &'static Self {
         let handler: &'static Self = Box::leak(Box::new(Self {
             api_keys,
             special_api_keys,
+            special_query_key_signers,
             api_key_payment_required,
             subscriptions,
+            subscription_domains,
             subscription_query_counters: RwLock::default(),
         }));
 
@@ -186,6 +188,27 @@ impl AuthHandler {
             AuthToken::ApiKey(_) => return Ok(()),
             AuthToken::Ticket(payload, subscription) => (payload, subscription),
         };
+
+        // This is safe, since we have already verified the signature and the claimed signer match.
+        // This is placed before the subscriptions domain check to allow the same special query keys to be used across
+        // testnet & mainnet.
+        if self
+            .special_query_key_signers
+            .contains(&Address::from(ticket_payload.signer.0))
+        {
+            return Ok(());
+        }
+
+        let matches_subscriptions_domain = self
+            .subscription_domains
+            .get(&ticket_payload.chain_id)
+            .map(|contract| contract == &Address::from(ticket_payload.contract.0))
+            .unwrap_or(false);
+        ensure!(
+            matches_subscriptions_domain,
+            "Query key chain_id or contract not allowed"
+        );
+
         let user: Address = ticket_payload.user().0.into();
         let counters = match self.subscription_query_counters.try_read() {
             Ok(counters) => counters,

graph-gateway/src/config.rs

@@ -172,10 +172,12 @@ impl FromStr for SignerKey {
 #[serde_as]
 #[derive(Debug, Deserialize)]
 pub struct Subscriptions {
-    /// Subscriptions contract owners
-    pub contract_owners: Vec<Address>,
+    /// Subscriptions contract domains
+    pub domains: Vec<SubscriptionsDomain>,
     /// Kafka topic to report subscription queries
     pub kafka_topic: Option<String>,
+    /// Query key signers that don't require payment
+    pub special_signers: Vec<Address>,
     /// Subscriptions subgraph URL
     #[serde_as(as = "DisplayFromStr")]
     pub subgraph: Url,
@@ -186,3 +188,10 @@ pub struct Subscriptions {
     #[serde_as(as = "FromInto<Vec<SubscriptionTier>>")]
     pub tiers: SubscriptionTiers,
 }
+
+#[serde_as]
+#[derive(Debug, Deserialize)]
+pub struct SubscriptionsDomain {
+    pub chain_id: u64,
+    pub contract: Address,
+}

graph-gateway/src/main.rs

@@ -213,23 +213,33 @@ async fn main() {
     let subscription_tiers: &'static SubscriptionTiers =
         Box::leak(Box::new(subscription_tiers.unwrap_or_default()));
 
-    let subscriptions = match config.subscriptions {
+    let subscriptions = match &config.subscriptions {
         None => Eventual::from_value(Ptr::default()),
         Some(subscriptions) => subscriptions_subgraph::Client::create(
             subgraph_client::Client::new(
                 http_client.clone(),
-                subscriptions.subgraph,
-                subscriptions.ticket,
+                subscriptions.subgraph.clone(),
+                subscriptions.ticket.clone(),
             ),
             subscription_tiers,
-            subscriptions.contract_owners,
         ),
     };
     let auth_handler = AuthHandler::create(
         api_keys,
         HashSet::from_iter(config.special_api_keys),
+        config
+            .subscriptions
+            .iter()
+            .flat_map(|s| s.special_signers.clone())
+            .collect(),
         config.api_key_payment_required,
         subscriptions,
+        config
+            .subscriptions
+            .iter()
+            .flat_map(|s| &s.domains)
+            .map(|d| (d.chain_id, d.contract))
+            .collect(),
     );
     let query_fees_target = config
         .query_fees_target

graph-gateway/src/subscriptions_subgraph.rs

@@ -13,34 +13,18 @@ use crate::subscriptions::{ActiveSubscription, Subscription};
 pub struct Client {
     subgraph_client: subgraph_client::Client,
     tiers: &'static SubscriptionTiers,
-    owner_subscriptions: Vec<(Address, Subscription)>,
     subscriptions: EventualWriter<Ptr<HashMap<Address, Subscription>>>,
 }
 
 impl Client {
     pub fn create(
         subgraph_client: subgraph_client::Client,
         tiers: &'static SubscriptionTiers,
-        contract_owners: Vec<Address>,
     ) -> Eventual<Ptr<HashMap<Address, Subscription>>> {
-        // TODO: query contract owners and authorized signers from subrgaph
-        let owner_subscriptions: Vec<(Address, Subscription)> = contract_owners
-            .iter()
-            .map(|owner| {
-                let sub = Subscription {
-                    queries_per_minute: u32::MAX,
-                    signers: vec![],
-                };
-                (*owner, sub)
-            })
-            .collect();
-
-        let (mut subscriptions_tx, subscriptions_rx) = Eventual::new();
-        subscriptions_tx.write(Ptr::new(owner_subscriptions.iter().cloned().collect()));
+        let (subscriptions_tx, subscriptions_rx) = Eventual::new();
         let client = Arc::new(Mutex::new(Client {
             subgraph_client,
             tiers,
-            owner_subscriptions,
             subscriptions: subscriptions_tx,
         }));
 
@@ -114,7 +98,6 @@ impl Client {
                 };
                 (user.id, sub)
             })
-            .chain(self.owner_subscriptions.clone())
             .collect();
         self.subscriptions.write(Ptr::new(subscriptions_map));