ENG-207: Restrict namespaces (#133)

* Restrict namespaces

* Add env var to restrict namespaces

Author: calummoore

helm/mainnet-values.yaml

@@ -21,6 +21,8 @@ service:
     cloud.google.com/neg: '{"ingress": true}'
 
 extraEnvVars:
+  - name: RESTRICT_NAMESPACES
+    value: "true"
   - name: PEERS
     value: "12D3KooWR9TtbDy2hebAU2qgbP2e3UGX8BrrZUDY6eiXHvQF3Nn4,12D3KooWKnx6YDNsZhMnqr9UADPLZXtAkYkDEo6Vs7KQjAL2dWCG,12D3KooWDtEwZyL2tjW96VCE5bJgDP58w31JTkeavaYafuA7oXcD"
   - name: WHITELIST

helm/prenet-values.yaml

@@ -20,6 +20,8 @@ persistence:
    size: 20Gi
 
 extraEnvVars:
+   - name: RESTRICT_NAMESPACES
+     value: "true"
    - name: PEERS
      value: "12D3KooWGwYHQGP9xJ8uxrmfhZMubfAsWPrTByGgq2N8hH2atMJZ,12D3KooWKLvUcUc4h3BEKGHKMz2kxfS7Xpps2kEfaQzpchGmMgSy,12D3KooWLYByKfin35Dh37YJf8HwQfyrL2beJW67712y9kChgbja"
    - name: DIAL_ADDR

helm/testnet-values.yaml

@@ -18,6 +18,8 @@ resources:
 additionalDomain: testnet.spacetime.xyz
 
 extraEnvVars:
+  - name: RESTRICT_NAMESPACES
+    value: "true"
   - name: PEERS
     value: "12D3KooWSZxfWBu726XtsGumTFtmqFTTenud1HempLgHx5kSvBaM,12D3KooWNSEqaQB4ZHFrdw4A1efz3G5aTxzfRPXPjrT5z2yfGaTs,12D3KooWPBFaDpnqiPRTQZqvf1W41fcszS27qPhpkL757Z2XSkcz"
   - name: SNAPSHOT_CHUNK_SIZE

polybase/src/config.rs

@@ -87,6 +87,10 @@ pub struct Config {
     /// Public key whitelist
     #[arg(long, env = "WHITELIST", value_parser, value_delimiter = ',')]
     pub whitelist: Option<Vec<String>>,
+
+    /// Restrict namespaces to pk/<pk>/<collection_name>
+    #[arg(long, env = "RESTRICT_NAMESPACES", default_value = "false")]
+    pub restrict_namespaces: bool,
 }
 
 #[derive(Subcommand, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)]

polybase/src/errors.rs

@@ -43,6 +43,14 @@ pub enum AppError {
     #[error("invalid request")]
     B58(#[from] bs58::decode::Error),
 
+    #[error("anonymous namespaces are not allowed, sign your request")]
+    AnonNamespace,
+
     #[error("public key not included in allowed whitelist")]
     Whitelist,
+
+    #[error(
+        "namespace is invalid, must be in format pk/<public_key_hex>/<CollectionName> got {0}"
+    )]
+    InvalidNamespace(String),
 }

polybase/src/main.rs

@@ -238,6 +238,7 @@ async fn main() -> Result<()> {
         config.rpc_laddr,
         Arc::clone(&db),
         Arc::new(config.whitelist.clone()),
+        Arc::new(config.restrict_namespaces),
         logger.clone(),
     )?;
 

polybase/src/rpc.rs

@@ -19,6 +19,7 @@ use std::{borrow::Cow, cmp::min, sync::Arc, time::Duration};
 struct RouteState {
     db: Arc<Db>,
     whitelist: Arc<Option<Vec<String>>>,
+    restrict_namespaces: Arc<bool>,
 }
 
 #[get("/")]
@@ -340,9 +341,15 @@ async fn post_record(
     let auth = body.auth.map(|a| a.into());
     let db: Arc<_> = Arc::clone(&state.db);
 
-    // Check whitelist
+    // New collection is being created
     if collection_id == "Collection" {
-        validate_whitelist(&state.whitelist, &auth)?;
+        // Validate whitelist (if it exists)
+        validate_new_collection(
+            &body.data.args[0],
+            &state.whitelist,
+            &state.restrict_namespaces,
+            &auth,
+        )?;
     }
 
     let txn = CallTxn::new(
@@ -437,6 +444,7 @@ pub fn create_rpc_server(
     rpc_laddr: String,
     db: Arc<Db>,
     whitelist: Arc<Option<Vec<String>>>,
+    restrict_namespaces: Arc<bool>,
     logger: slog::Logger,
 ) -> Result<Server, std::io::Error> {
     Ok(HttpServer::new(move || {
@@ -446,6 +454,7 @@ pub fn create_rpc_server(
             .app_data(web::Data::new(RouteState {
                 db: Arc::clone(&db),
                 whitelist: Arc::clone(&whitelist),
+                restrict_namespaces: Arc::clone(&restrict_namespaces),
             }))
             .wrap(SlogMiddleware::new(logger.clone()))
             .wrap(cors)
@@ -464,27 +473,57 @@ pub fn create_rpc_server(
     .run())
 }
 
-fn validate_whitelist(
+fn validate_new_collection(
+    collection_id: &serde_json::Value,
     whitelist: &Option<Vec<String>>,
+    restrict_namespaces: &bool,
     auth: &Option<AuthUser>,
 ) -> Result<(), HTTPError> {
-    // Check whitelist
+    let pk = auth
+        .as_ref()
+        .map(|a| a.public_key().to_hex().unwrap_or("".to_string()))
+        .unwrap_or("".to_string());
+
+    // Check collection whitelist
     if let Some(whitelist) = whitelist {
-        if let Some(auth_user) = auth {
-            // Convert the key to hex for easier comparison
-            let pk = auth_user.public_key().to_hex().unwrap_or("".to_string());
-            if pk.is_empty() || !whitelist.contains(&pk) {
-                return Err(HTTPError::new(
-                    ReasonCode::Unauthorized,
-                    Some(Box::new(AppError::Whitelist)),
-                ));
-            }
-        } else {
+        if pk.is_empty() || !whitelist.contains(&pk) {
             return Err(HTTPError::new(
                 ReasonCode::Unauthorized,
                 Some(Box::new(AppError::Whitelist)),
             ));
         }
     }
+
+    // Check namespace is valid (only pk/<pk> currently allowed)
+    if *restrict_namespaces {
+        if pk.is_empty() {
+            return Err(HTTPError::new(
+                ReasonCode::Unauthorized,
+                Some(Box::new(AppError::AnonNamespace)),
+            ));
+        }
+
+        match collection_id {
+            serde_json::Value::String(id) => {
+                let parts: Vec<&str> = id.split('/').collect();
+                if parts.len() != 3 || parts[0] != "pk" || parts[1] != pk {
+                    return Err(HTTPError::new(
+                        ReasonCode::Unauthorized,
+                        Some(Box::new(AppError::InvalidNamespace(id.clone()))),
+                    ));
+                }
+            }
+            _ => {
+                return Err(HTTPError::new(
+                    ReasonCode::Unauthorized,
+                    Some(Box::new(AppError::InvalidNamespace(format!(
+                        "{:?}",
+                        collection_id
+                    )))),
+                ));
+            }
+        }
+    }
+
     Ok(())
 }

polybase/tests/api/mod.rs

@@ -11,6 +11,7 @@ mod index_where_sort;
 mod map_field;
 mod nested_field;
 mod other_collection_fns;
+mod restrict_namespaces;
 mod start_stop;
 mod store_other_collection_records;
 mod whitelist;
@@ -104,6 +105,7 @@ impl PortPool {
 struct ServerConfig {
     whitelist: Option<Vec<String>>,
     keep_port_after_drop: bool,
+    restrict_namespaces: bool,
 }
 
 #[derive(Debug)]
@@ -138,6 +140,10 @@ impl Server {
             if let Some(ref whitelist) = config.whitelist {
                 command.arg("--whitelist").arg(whitelist.join(","));
             }
+
+            if config.restrict_namespaces {
+                command.arg("--restrict-namespaces");
+            }
         }
 
         command.arg("--root-dir").arg(root_dir.path());

polybase/tests/api/restrict_namespaces.rs

@@ -0,0 +1,168 @@
+use crate::api::{Error, ErrorData, Signature, Signer};
+use std::time::SystemTime;
+
+use super::{Server, ServerConfig};
+
+#[derive(Debug, PartialEq, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct Account {
+    id: String,
+}
+
+#[tokio::test]
+async fn valid() {
+    let schema = r#"
+@public
+collection Account {
+    id: string;
+}
+    "#;
+
+    let (private_key, public_key) = secp256k1::generate_keypair(&mut rand::thread_rng());
+    let public_key = indexer::PublicKey::from_secp256k1_key(&public_key).unwrap();
+    let pk_hex = public_key.to_hex().unwrap();
+
+    let signer = Signer::from(move |body: &str| {
+        let mut signature = Signature::create(&private_key, SystemTime::now(), body);
+        signature.public_key = Some(public_key.clone());
+        signature
+    });
+
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        restrict_namespaces: true,
+        ..Default::default()
+    }))
+    .await;
+
+    let collection_id = format!("pk/{}/Account", pk_hex);
+    let res = server
+        .create_collection::<Account>(collection_id.as_str(), schema, Some(&signer))
+        .await
+        .unwrap();
+
+    assert_eq!(res.id, collection_id.to_string());
+}
+
+#[tokio::test]
+async fn invalid_prefix() {
+    let schema = r#"
+@public
+collection Account {
+    id: string;
+}
+    "#;
+
+    let (private_key, public_key) = secp256k1::generate_keypair(&mut rand::thread_rng());
+    let public_key = indexer::PublicKey::from_secp256k1_key(&public_key).unwrap();
+    let pk_hex = public_key.to_hex().unwrap();
+
+    let signer = Signer::from(move |body: &str| {
+        let mut signature = Signature::create(&private_key, SystemTime::now(), body);
+        signature.public_key = Some(public_key.clone());
+        signature
+    });
+
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        restrict_namespaces: true,
+        ..Default::default()
+    }))
+    .await;
+
+    let collection_id = format!("other/{}/Account", pk_hex);
+
+    let res = server
+        .create_collection::<Account>(collection_id.as_str(), schema, Some(&signer))
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        res,
+        Error {
+            error: ErrorData {
+                code: "permission-denied".to_string(),
+                message: format!("namespace is invalid, must be in format pk/<public_key_hex>/<CollectionName> got {}", collection_id.as_str()),
+                reason: "unauthorized".to_string(),
+            }
+        }
+    );
+}
+
+#[tokio::test]
+async fn mismatch_pk() {
+    let schema = r#"
+@public
+collection Account {
+    id: string;
+}
+    "#;
+
+    // Key signing the request
+    let (private_key, public_key) = secp256k1::generate_keypair(&mut rand::thread_rng());
+    let public_key = indexer::PublicKey::from_secp256k1_key(&public_key).unwrap();
+
+    // Key to be used in whitelist
+    let (_, alt_public_key) = secp256k1::generate_keypair(&mut rand::thread_rng());
+    let alt_public_key = indexer::PublicKey::from_secp256k1_key(&alt_public_key).unwrap();
+    let pk_hex: String = alt_public_key.to_hex().unwrap();
+
+    let signer = Signer::from(move |body: &str| {
+        let mut signature = Signature::create(&private_key, SystemTime::now(), body);
+        signature.public_key = Some(public_key.clone());
+        signature
+    });
+
+    let collection_id = format!("pk/{}/Account", pk_hex);
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        restrict_namespaces: true,
+        ..Default::default()
+    }))
+    .await;
+
+    let res = server
+        .create_collection::<Account>(collection_id.as_str(), schema, Some(&signer))
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        res,
+        Error {
+            error: ErrorData {
+                code: "permission-denied".to_string(),
+                message: format!("namespace is invalid, must be in format pk/<public_key_hex>/<CollectionName> got {}", collection_id.as_str()),
+                reason: "unauthorized".to_string(),
+            }
+        }
+    );
+}
+
+#[tokio::test]
+async fn missing_pk() {
+    let schema = r#"
+@public
+collection Account {
+    id: string;
+}
+    "#;
+
+    let server = Server::setup_and_wait(Some(ServerConfig {
+        restrict_namespaces: true,
+        ..ServerConfig::default()
+    }))
+    .await;
+
+    let res = server
+        .create_collection::<Account>("test/Account", schema, None)
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        res,
+        Error {
+            error: ErrorData {
+                code: "permission-denied".to_string(),
+                message: "anonymous namespaces are not allowed, sign your request".to_string(),
+                reason: "unauthorized".to_string(),
+            }
+        }
+    );
+}