[AUDIT][MEDIUM][O-1] Orchestrator slots may be trivially exhausted

crates/orchestrator/run-config.toml

@@ -39,22 +39,19 @@ seed = [
 ]
 start_delay_seconds = 0
 cdn_marshal_address = "127.0.0.1:9000"
+public_keys = [
+    "BLS_VER_KEY~p-JKk1VvO1RoMrDrqyjz0P1VGwtOaEjF5jLjpOZbJi5O747fvYEOg0OvCl_CLe4shh7vsqeG9uMF9RssM12sLSuaiVJkCClxEI5mRLV4qff1UjZAZJIBgeL1_hRhRUkpqC0Trm1qtvXtZ8FwOCIzYXv8c300Au824k7FxjjcWLBL",
+]
 
 [config]
 num_nodes_with_stake = 10
 num_nodes_without_stake = 0
-start_threshold = [
-    8,
-    10,
-]
+start_threshold = [8, 10]
 staked_da_nodes = 10
 non_staked_da_nodes = 0
 fixed_leader_for_gpuvid = 1
 next_view_timeout = 30000
-timeout_ratio = [
-    11,
-    10,
-]
+timeout_ratio = [11, 10]
 round_start_delay = 1
 start_delay = 1
 num_bootstrap = 5

crates/orchestrator/src/config.rs

@@ -208,6 +208,10 @@ pub struct NetworkConfig<KEY: SignatureKey> {
     pub builder: BuilderType,
     /// random builder config
     pub random_builder: Option<RandomBuilderConfig>,
+    /// The list of public keys that are allowed to connect to the orchestrator
+    pub public_keys: Vec<KEY>,
+    /// Whether or not to disable registration verification.
+    pub disable_registration_verification: bool,
 }
 
 /// the source of the network config
@@ -439,6 +443,8 @@ impl<K: SignatureKey> Default for NetworkConfig<K> {
             commit_sha: String::new(),
             builder: BuilderType::default(),
             random_builder: None,
+            public_keys: vec![],
+            disable_registration_verification: false,
         }
     }
 }
@@ -491,10 +497,23 @@ pub struct NetworkConfigFile<KEY: SignatureKey> {
     /// random builder configuration
     #[serde(default)]
     pub random_builder: Option<RandomBuilderConfig>,
+    /// The list of public keys that are allowed to connect to the orchestrator
+    #[serde(default)]
+    pub public_keys: Vec<KEY>,
+    /// Whether or not to disable registration verification.
+    #[serde(default)]
+    pub disable_registration_verification: bool,
 }
 
 impl<K: SignatureKey> From<NetworkConfigFile<K>> for NetworkConfig<K> {
     fn from(val: NetworkConfigFile<K>) -> Self {
+        #[cfg(not(debug_assertions))]
+        {
+            if val.disable_registration_verification {
+                panic!("Registration verification cannot be turned off in production builds");
+            }
+        }
+
         NetworkConfig {
             rounds: val.rounds,
             indexed_da: val.indexed_da,
@@ -536,6 +555,8 @@ impl<K: SignatureKey> From<NetworkConfigFile<K>> for NetworkConfig<K> {
             commit_sha: String::new(),
             builder: val.builder,
             random_builder: val.random_builder,
+            public_keys: val.public_keys,
+            disable_registration_verification: val.disable_registration_verification,
         }
     }
 }

crates/orchestrator/src/lib.rs

@@ -12,7 +12,7 @@ pub mod client;
 pub mod config;
 
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     fs::OpenOptions,
     io::{self, ErrorKind},
     time::Duration,
@@ -23,7 +23,10 @@ use client::{BenchResults, BenchResultsDownloadConfig};
 use config::BuilderType;
 use csv::Writer;
 use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
-use hotshot_types::{traits::signature_key::SignatureKey, PeerConfig};
+use hotshot_types::{
+    traits::signature_key::{SignatureKey, StakeTableEntryType},
+    PeerConfig,
+};
 use libp2p::{
     identity::{
         ed25519::{Keypair as EdKeypair, SecretKey},
@@ -83,6 +86,8 @@ struct OrchestratorState<KEY: SignatureKey> {
     peer_pub_ready: bool,
     /// A map from public keys to `(node_index, is_da)`.
     pub_posted: HashMap<Vec<u8>, (u64, bool)>,
+    /// A collection for storing `ready` public keys.
+    ready_posted: HashSet<KEY>,
     /// Whether nodes should start their HotShot instances
     /// Will be set to true once all nodes post they are ready to start
     start: bool,
@@ -114,6 +119,7 @@ impl<KEY: SignatureKey + 'static> OrchestratorState<KEY> {
             config: network_config,
             peer_pub_ready: false,
             pub_posted: HashMap::new(),
+            ready_posted: HashSet::new(),
             nodes_connected: 0,
             start: false,
             bench_results: BenchResults::default(),
@@ -228,7 +234,7 @@ pub trait OrchestratorApi<KEY: SignatureKey> {
     /// post endpoint for whether or not all nodes are ready
     /// # Errors
     /// if unable to serve
-    fn post_ready(&mut self) -> Result<(), ServerError>;
+    fn post_ready(&mut self, pubkey: &mut Vec<u8>) -> Result<(), ServerError>;
     /// post endpoint for manually starting the orchestrator
     /// # Errors
     /// if unable to serve
@@ -328,7 +334,21 @@ where
 
         let node_index = self.pub_posted.len() as u64;
 
+        // Deserialize the public key
         let staked_pubkey = PeerConfig::<KEY>::from_bytes(pubkey).unwrap();
+
+        // Check if the node is allowed to connect
+        if !self
+            .config
+            .public_keys
+            .contains(&staked_pubkey.stake_table_entry.public_key())
+        {
+            return Err(ServerError {
+                status: tide_disco::StatusCode::FORBIDDEN,
+                message: "You are unauthorized to register with the orchestrator".to_string(),
+            });
+        }
+
         self.config
             .config
             .known_nodes_with_stake
@@ -416,8 +436,33 @@ where
     }
 
     // Assumes nodes do not post 'ready' twice
-    // TODO ED Add a map to verify which nodes have posted they're ready
-    fn post_ready(&mut self) -> Result<(), ServerError> {
+    fn post_ready(&mut self, pubkey: &mut Vec<u8>) -> Result<(), ServerError> {
+        // Deserialize the public key
+        let staked_pubkey = PeerConfig::<KEY>::from_bytes(pubkey).unwrap();
+
+        // Is this node allowed to connect?
+        if !self
+            .config
+            .public_keys
+            .contains(&staked_pubkey.stake_table_entry.public_key())
+        {
+            return Err(ServerError {
+                status: tide_disco::StatusCode::FORBIDDEN,
+                message: "You are unauthorized to register with the orchestrator".to_string(),
+            });
+        }
+
+        // Have they already connected?
+        if !self
+            .ready_posted
+            .insert(staked_pubkey.stake_table_entry.public_key().clone())
+        {
+            return Err(ServerError {
+                status: tide_disco::StatusCode::BAD_REQUEST,
+                message: "You have already reported your ready status".to_string(),
+            });
+        }
+
         self.nodes_connected += 1;
 
         println!("Nodes connected: {}", self.nodes_connected);
@@ -636,7 +681,23 @@ where
     })?
     .post(
         "post_ready",
-        |_req, state: &mut <State as ReadState>::State| async move { state.post_ready() }.boxed(),
+        |req, state: &mut <State as ReadState>::State| {
+            async move {
+                let mut body_bytes = req.body_bytes();
+                body_bytes.drain(..12);
+                // Decode the payload-supplied pubkey
+                let Ok(mut pubkey) =
+                    vbs::Serializer::<OrchestratorVersion>::deserialize(&body_bytes)
+                else {
+                    return Err(ServerError {
+                        status: tide_disco::StatusCode::BAD_REQUEST,
+                        message: "Malformed body".to_string(),
+                    });
+                };
+                state.post_ready(&mut pubkey)
+            }
+            .boxed()
+        },
     )?
     .post(
         "post_manual_start",

crates/types/src/stake_table.rs

@@ -21,11 +21,16 @@ pub struct StakeTableEntry<K: SignatureKey> {
     pub stake_amount: U256,
 }
 
-impl<K: SignatureKey> StakeTableEntryType for StakeTableEntry<K> {
+impl<K: SignatureKey> StakeTableEntryType<K> for StakeTableEntry<K> {
     /// Get the stake amount
     fn stake(&self) -> U256 {
         self.stake_amount
     }
+
+    /// Get the public key
+    fn public_key(&self) -> K {
+        self.stake_key.clone()
+    }
 }
 
 impl<K: SignatureKey> StakeTableEntry<K> {

crates/types/src/traits/signature_key.rs

@@ -25,9 +25,11 @@ use super::EncodeBytes;
 use crate::{utils::BuilderCommitment, vid::VidSchemeType};
 
 /// Type representing stake table entries in a `StakeTable`
-pub trait StakeTableEntryType {
+pub trait StakeTableEntryType<K> {
     /// Get the stake value
     fn stake(&self) -> U256;
+    /// Get the public key
+    fn public_key(&self) -> K;
 }
 
 /// Trait for abstracting public key signatures
@@ -60,7 +62,7 @@ pub trait SignatureKey:
         + for<'a> Deserialize<'a>
         + Hash;
     /// The type of the entry that contain both public key and stake value
-    type StakeTableEntry: StakeTableEntryType
+    type StakeTableEntry: StakeTableEntryType<Self>
         + Send
         + Sync
         + Sized