Skip to main content

hotshot_orchestrator/
lib.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7//! Orchestrator for manipulating nodes and recording results during a run of `HotShot` tests
8
9/// The orchestrator's clients
10pub mod client;
11
12use std::{
13    collections::{HashMap, HashSet},
14    fs,
15    fs::OpenOptions,
16    io,
17    sync::Arc,
18    time::Duration,
19};
20
21use alloy::primitives::U256;
22use async_lock::RwLock;
23use axum::{
24    Router,
25    body::Bytes,
26    extract::{Path, State},
27    http::{HeaderMap, StatusCode},
28    response::Response,
29    routing::{get, post},
30};
31use client::{BenchResults, BenchResultsDownloadConfig};
32use csv::Writer;
33use futures::{StreamExt, stream::FuturesUnordered};
34use hotshot_types::{
35    PeerConfig,
36    network::{BuilderType, NetworkConfig, PublicKeysFile},
37    traits::{
38        node_implementation::NodeType,
39        signature_key::{SignatureKey, StakeTableEntryType},
40    },
41};
42use http_client::{Url, error::ClientErr};
43use http_wire::{self as wire, ServerError, WireVersion, cors_layer, healthcheck_response};
44use serde::de::DeserializeOwned;
45use tokio::net::TcpListener;
46use vbs::{BinarySerializer, Serializer};
47
48/// VBS framing version for orchestrator request bodies and the client in `client.rs`.
49///
50/// Responses go through [`wire::respond`], which always encodes at [`WireVersion`]; aliasing
51/// keeps request decoding and the client in lockstep with the responses.
52pub type OrchestratorVersion = WireVersion;
53
54/// The state of the orchestrator
55#[derive(Default, Clone)]
56#[allow(clippy::struct_excessive_bools)]
57struct OrchestratorState<TYPES: NodeType> {
58    /// Tracks the latest node index we have generated a configuration for
59    latest_index: u16,
60    /// Tracks the latest temporary index we have generated for init validator's key pair
61    tmp_latest_index: u16,
62    /// The network configuration
63    config: NetworkConfig<TYPES>,
64    /// Whether the network configuration has been updated with all the peer's public keys/configs
65    peer_pub_ready: bool,
66    /// A map from public keys to `(node_index, is_da)`.
67    pub_posted: HashMap<Vec<u8>, (u64, bool)>,
68    /// Whether nodes should start their HotShot instances
69    /// Will be set to true once all nodes post they are ready to start
70    start: bool,
71    /// The total nodes that have posted they are ready to start
72    nodes_connected: HashSet<PeerConfig<TYPES>>,
73    /// The results of the benchmarks
74    bench_results: BenchResults,
75    /// The number of nodes that have posted their results
76    nodes_post_results: u64,
77    /// Whether the orchestrator can be started manually
78    manual_start_allowed: bool,
79    /// Whether we are still accepting new keys for registration
80    accepting_new_keys: bool,
81    /// Builder address pool
82    builders: Vec<Url>,
83    /// whether we are using a fixed stake table, disabling public key registration
84    fixed_stake_table: bool,
85}
86
87impl<TYPES: NodeType> OrchestratorState<TYPES> {
88    /// create a new [`OrchestratorState`]
89    pub fn new(network_config: NetworkConfig<TYPES>) -> Self {
90        let mut peer_pub_ready = false;
91        let mut fixed_stake_table = false;
92
93        if network_config.config.known_nodes_with_stake.is_empty() {
94            println!(
95                "No nodes were loaded from the config file. Nodes will be allowed to register \
96                 dynamically."
97            );
98        } else {
99            println!("Initializing orchestrator with fixed stake table.");
100            peer_pub_ready = true;
101            fixed_stake_table = true;
102        }
103
104        let builders = if matches!(network_config.builder, BuilderType::External) {
105            network_config.config.builder_urls.clone().into()
106        } else {
107            vec![]
108        };
109
110        OrchestratorState {
111            latest_index: 0,
112            tmp_latest_index: 0,
113            config: network_config,
114            peer_pub_ready,
115            pub_posted: HashMap::new(),
116            nodes_connected: HashSet::new(),
117            start: false,
118            bench_results: BenchResults::default(),
119            nodes_post_results: 0,
120            manual_start_allowed: true,
121            accepting_new_keys: true,
122            builders,
123            fixed_stake_table,
124        }
125    }
126
127    /// Output the results to a csv file according to orchestrator state
128    pub fn output_to_csv(&self) {
129        let output_csv = BenchResultsDownloadConfig {
130            commit_sha: self.config.commit_sha.clone(),
131            total_nodes: self.config.config.num_nodes_with_stake.into(),
132            da_committee_size: self.config.config.da_staked_committee_size,
133            fixed_leader_for_gpuvid: self.config.config.fixed_leader_for_gpuvid,
134            transactions_per_round: self.config.transactions_per_round,
135            transaction_size: self.bench_results.transaction_size_in_bytes,
136            rounds: self.config.rounds,
137            partial_results: self.bench_results.partial_results.clone(),
138            avg_latency_in_sec: self.bench_results.avg_latency_in_sec,
139            minimum_latency_in_sec: self.bench_results.minimum_latency_in_sec,
140            maximum_latency_in_sec: self.bench_results.maximum_latency_in_sec,
141            throughput_bytes_per_sec: self.bench_results.throughput_bytes_per_sec,
142            total_transactions_committed: self.bench_results.total_transactions_committed,
143            total_time_elapsed_in_sec: self.bench_results.total_time_elapsed_in_sec,
144            total_num_views: self.bench_results.total_num_views,
145            failed_num_views: self.bench_results.failed_num_views,
146            committee_type: self.bench_results.committee_type.clone(),
147        };
148        // Open the CSV file in append mode
149        let results_csv_file = OpenOptions::new()
150            .create(true)
151            .append(true) // Open in append mode
152            .open("scripts/benchmarks_results/results.csv")
153            .unwrap();
154        // Open a file for writing
155        let mut wtr = Writer::from_writer(results_csv_file);
156        let _ = wtr.serialize(output_csv);
157        let _ = wtr.flush();
158        println!("Results successfully saved in scripts/benchmarks_results/results.csv");
159    }
160}
161
162/// An api exposed by the orchestrator
163pub trait OrchestratorApi<TYPES: NodeType> {
164    /// Post an identity to the orchestrator.
165    /// # Errors
166    /// If we were unable to serve the request
167    fn post_identity(&mut self) -> Result<u16, ServerError>;
168    /// post endpoint for each node's config
169    /// # Errors
170    /// if unable to serve
171    fn post_getconfig(&mut self, _node_index: u16) -> Result<NetworkConfig<TYPES>, ServerError>;
172    /// get endpoint for the next available temporary node index
173    /// # Errors
174    /// if unable to serve
175    fn get_tmp_node_index(&mut self) -> Result<u16, ServerError>;
176    /// post endpoint for each node's public key
177    /// # Errors
178    /// if unable to serve
179    fn register_public_key(
180        &mut self,
181        pubkey: &mut Vec<u8>,
182        is_da: bool,
183    ) -> Result<(u64, bool), ServerError>;
184    /// post endpoint for whether or not all peers public keys are ready
185    /// # Errors
186    /// if unable to serve
187    fn peer_pub_ready(&self) -> Result<bool, ServerError>;
188    /// get endpoint for the network config after all peers public keys are collected
189    /// # Errors
190    /// if unable to serve
191    fn post_config_after_peer_collected(&mut self) -> Result<NetworkConfig<TYPES>, ServerError>;
192    /// get endpoint for whether or not the run has started
193    /// # Errors
194    /// if unable to serve
195    fn get_start(&self) -> Result<bool, ServerError>;
196    /// post endpoint for the results of the run
197    /// # Errors
198    /// if unable to serve
199    fn post_run_results(&mut self, metrics: BenchResults) -> Result<(), ServerError>;
200    /// A node POSTs its public key to let the orchestrator know that it is ready
201    /// # Errors
202    /// if unable to serve
203    fn post_ready(&mut self, peer_config: &PeerConfig<TYPES>) -> Result<(), ServerError>;
204    /// post endpoint for manually starting the orchestrator
205    /// # Errors
206    /// if unable to serve
207    fn post_manual_start(&mut self, password_bytes: Vec<u8>) -> Result<(), ServerError>;
208    /// post endpoint for registering a builder with the orchestrator
209    /// # Errors
210    /// if unable to serve
211    fn post_builder(&mut self, builder: Url) -> Result<(), ServerError>;
212    /// get endpoints for builders
213    /// # Errors
214    /// if not all builders are registered yet
215    fn get_builders(&self) -> Result<Vec<Url>, ServerError>;
216}
217
218impl<TYPES: NodeType> OrchestratorState<TYPES>
219where
220    TYPES::SignatureKey: serde::Serialize + Clone + SignatureKey + 'static,
221{
222    /// register a node with an unknown public key.
223    /// this method should be used when we don't have a fixed stake table
224    fn register_unknown(
225        &mut self,
226        pubkey: &mut Vec<u8>,
227        da_requested: bool,
228    ) -> Result<(u64, bool), ServerError> {
229        if let Some((node_index, is_da)) = self.pub_posted.get(pubkey) {
230            return Ok((*node_index, *is_da));
231        }
232
233        if !self.accepting_new_keys {
234            return Err(ServerError {
235                status: StatusCode::FORBIDDEN,
236                message: "Network has been started manually, and is no longer registering new \
237                          keys."
238                    .to_string(),
239            });
240        }
241
242        let node_index = self.pub_posted.len() as u64;
243
244        // Deserialize the public key
245        let staked_pubkey = PeerConfig::<TYPES>::from_bytes(pubkey).unwrap();
246
247        self.config
248            .config
249            .known_nodes_with_stake
250            .push(staked_pubkey.clone());
251
252        let mut added_to_da = false;
253
254        let da_full =
255            self.config.config.known_da_nodes.len() >= self.config.config.da_staked_committee_size;
256
257        #[allow(clippy::nonminimal_bool)]
258        // We add the node to the DA committee depending on either its node index or whether it requested membership.
259        //
260        // Since we issue `node_index` incrementally, if we are deciding DA membership by node_index
261        // we only need to check that the DA committee is not yet full.
262        //
263        // Note: this logically simplifies to (self.config.indexed_da || da_requested) && !da_full,
264        // but writing it that way makes it a little less clear to me.
265        if (self.config.indexed_da || (!self.config.indexed_da && da_requested)) && !da_full {
266            self.config.config.known_da_nodes.push(staked_pubkey);
267            added_to_da = true;
268        }
269
270        self.pub_posted
271            .insert(pubkey.clone(), (node_index, added_to_da));
272
273        tracing::error!("Posted public key for node_index {node_index}");
274
275        // node_index starts at 0, so once it matches `num_nodes_with_stake`
276        // we will have registered one node too many. hence, we want `node_index + 1`.
277        if node_index + 1 >= (self.config.config.num_nodes_with_stake.get() as u64) {
278            self.peer_pub_ready = true;
279            self.accepting_new_keys = false;
280        }
281        Ok((node_index, added_to_da))
282    }
283
284    /// register a node on the fixed stake table, which was loaded at startup
285    fn register_from_list(
286        &mut self,
287        pubkey: &mut Vec<u8>,
288        da_requested: bool,
289    ) -> Result<(u64, bool), ServerError> {
290        // if we've already registered this node before, we just retrieve its info from `pub_posted`
291        if let Some((node_index, is_da)) = self.pub_posted.get(pubkey) {
292            return Ok((*node_index, *is_da));
293        }
294
295        // Deserialize the public key
296        let staked_pubkey = PeerConfig::<TYPES>::from_bytes(pubkey).unwrap();
297
298        // Check if the node is allowed to connect, returning its index and config entry if so.
299        let Some((node_index, node_config)) =
300            self.config.public_keys.iter().enumerate().find(|keys| {
301                keys.1.stake_table_key == staked_pubkey.stake_table_entry.public_key()
302            })
303        else {
304            return Err(ServerError {
305                status: StatusCode::FORBIDDEN,
306                message: "You are unauthorized to register with the orchestrator".to_string(),
307            });
308        };
309
310        // Check that our recorded DA status for the node matches what the node actually requested
311        if node_config.da != da_requested {
312            return Err(ServerError {
313                status: StatusCode::BAD_REQUEST,
314                message: format!(
315                    "Mismatch in DA status in registration for node {}. DA requested: {}, \
316                     expected: {}",
317                    node_index, da_requested, node_config.da
318                ),
319            });
320        }
321
322        let added_to_da = node_config.da;
323
324        self.pub_posted
325            .insert(pubkey.clone(), (node_index as u64, added_to_da));
326
327        tracing::error!("Node {node_index} has registered.");
328
329        Ok((node_index as u64, added_to_da))
330    }
331}
332
333impl<TYPES: NodeType> OrchestratorApi<TYPES> for OrchestratorState<TYPES>
334where
335    TYPES::SignatureKey: serde::Serialize + Clone + SignatureKey + 'static,
336{
337    fn post_identity(&mut self) -> Result<u16, ServerError> {
338        let node_index = self.latest_index;
339        self.latest_index += 1;
340
341        if usize::from(node_index) >= self.config.config.num_nodes_with_stake.get() {
342            return Err(ServerError {
343                status: StatusCode::BAD_REQUEST,
344                message: "Network has reached capacity".to_string(),
345            });
346        }
347        Ok(node_index)
348    }
349
350    // Assumes nodes will set their own index that they received from the
351    // 'identity' endpoint
352    fn post_getconfig(&mut self, _node_index: u16) -> Result<NetworkConfig<TYPES>, ServerError> {
353        Ok(self.config.clone())
354    }
355
356    // Assumes one node do not get twice
357    fn get_tmp_node_index(&mut self) -> Result<u16, ServerError> {
358        let tmp_node_index = self.tmp_latest_index;
359        self.tmp_latest_index += 1;
360
361        if usize::from(tmp_node_index) >= self.config.config.num_nodes_with_stake.get() {
362            return Err(ServerError {
363                status: StatusCode::BAD_REQUEST,
364                message: "Node index getter for key pair generation has reached capacity"
365                    .to_string(),
366            });
367        }
368        Ok(tmp_node_index)
369    }
370
371    fn register_public_key(
372        &mut self,
373        pubkey: &mut Vec<u8>,
374        da_requested: bool,
375    ) -> Result<(u64, bool), ServerError> {
376        if self.fixed_stake_table {
377            self.register_from_list(pubkey, da_requested)
378        } else {
379            self.register_unknown(pubkey, da_requested)
380        }
381    }
382
383    fn peer_pub_ready(&self) -> Result<bool, ServerError> {
384        if !self.peer_pub_ready {
385            return Err(ServerError {
386                status: StatusCode::BAD_REQUEST,
387                message: "Peer's public configs are not ready".to_string(),
388            });
389        }
390        Ok(self.peer_pub_ready)
391    }
392
393    fn post_config_after_peer_collected(&mut self) -> Result<NetworkConfig<TYPES>, ServerError> {
394        if !self.peer_pub_ready {
395            return Err(ServerError {
396                status: StatusCode::BAD_REQUEST,
397                message: "Peer's public configs are not ready".to_string(),
398            });
399        }
400
401        Ok(self.config.clone())
402    }
403
404    fn get_start(&self) -> Result<bool, ServerError> {
405        // println!("{}", self.start);
406        if !self.start {
407            return Err(ServerError {
408                status: StatusCode::BAD_REQUEST,
409                message: "Network is not ready to start".to_string(),
410            });
411        }
412        Ok(self.start)
413    }
414
415    // Assumes nodes do not post 'ready' twice
416    fn post_ready(&mut self, peer_config: &PeerConfig<TYPES>) -> Result<(), ServerError> {
417        // If we have not disabled registration verification.
418        // Is this node allowed to connect?
419        if !self
420            .config
421            .config
422            .known_nodes_with_stake
423            .contains(peer_config)
424        {
425            return Err(ServerError {
426                status: StatusCode::FORBIDDEN,
427                message: "You are unauthorized to register with the orchestrator".to_string(),
428            });
429        }
430
431        // `HashSet::insert()` returns whether the node was newly inserted (true) or not
432        if self.nodes_connected.insert(peer_config.clone()) {
433            tracing::error!(
434                "Node {peer_config} connected. Total nodes connected: {}",
435                self.nodes_connected.len()
436            );
437        }
438
439        // i.e. nodes_connected >= num_nodes_with_stake * (start_threshold.0 / start_threshold.1)
440        if self.nodes_connected.len() as u64 * self.config.config.start_threshold.1
441            >= (self.config.config.num_nodes_with_stake.get() as u64)
442                * self.config.config.start_threshold.0
443        {
444            self.accepting_new_keys = false;
445            self.manual_start_allowed = false;
446            self.start = true;
447        }
448
449        Ok(())
450    }
451
452    /// Manually start the network
453    fn post_manual_start(&mut self, password_bytes: Vec<u8>) -> Result<(), ServerError> {
454        if !self.manual_start_allowed {
455            return Err(ServerError {
456                status: StatusCode::FORBIDDEN,
457                message: "Configs have already been distributed to nodes, and the network can no \
458                          longer be started manually."
459                    .to_string(),
460            });
461        }
462
463        let password = String::from_utf8(password_bytes)
464            .expect("Failed to decode raw password as UTF-8 string.");
465
466        // Check that the password matches
467        if self.config.manual_start_password != Some(password) {
468            return Err(ServerError {
469                status: StatusCode::FORBIDDEN,
470                message: "Incorrect password.".to_string(),
471            });
472        }
473
474        let registered_nodes_with_stake = self.config.config.known_nodes_with_stake.len();
475        let registered_da_nodes = self.config.config.known_da_nodes.len();
476
477        if registered_da_nodes > 1 {
478            self.config.config.num_nodes_with_stake =
479                std::num::NonZeroUsize::new(registered_nodes_with_stake)
480                    .expect("Failed to convert to NonZeroUsize; this should be impossible.");
481
482            self.config.config.da_staked_committee_size = registered_da_nodes;
483        } else {
484            return Err(ServerError {
485                status: StatusCode::FORBIDDEN,
486                message: format!(
487                    "We cannot manually start the network, because we only have \
488                     {registered_nodes_with_stake} nodes with stake registered, with \
489                     {registered_da_nodes} DA nodes."
490                ),
491            });
492        }
493
494        self.accepting_new_keys = false;
495        self.manual_start_allowed = false;
496        self.peer_pub_ready = true;
497        self.start = true;
498
499        Ok(())
500    }
501
502    // Aggregates results of the run from all nodes
503    fn post_run_results(&mut self, metrics: BenchResults) -> Result<(), ServerError> {
504        if metrics.total_transactions_committed != 0 {
505            // Deal with the bench results
506            if self.bench_results.total_transactions_committed == 0 {
507                self.bench_results = metrics;
508            } else {
509                // Deal with the bench results from different nodes
510                let cur_metrics = self.bench_results.clone();
511                self.bench_results.avg_latency_in_sec = (metrics.avg_latency_in_sec
512                    * metrics.num_latency
513                    + cur_metrics.avg_latency_in_sec * cur_metrics.num_latency)
514                    / (metrics.num_latency + cur_metrics.num_latency);
515                self.bench_results.num_latency += metrics.num_latency;
516                self.bench_results.minimum_latency_in_sec = metrics
517                    .minimum_latency_in_sec
518                    .min(cur_metrics.minimum_latency_in_sec);
519                self.bench_results.maximum_latency_in_sec = metrics
520                    .maximum_latency_in_sec
521                    .max(cur_metrics.maximum_latency_in_sec);
522                self.bench_results.throughput_bytes_per_sec = metrics
523                    .throughput_bytes_per_sec
524                    .max(cur_metrics.throughput_bytes_per_sec);
525                self.bench_results.total_transactions_committed = metrics
526                    .total_transactions_committed
527                    .max(cur_metrics.total_transactions_committed);
528                self.bench_results.total_time_elapsed_in_sec = metrics
529                    .total_time_elapsed_in_sec
530                    .max(cur_metrics.total_time_elapsed_in_sec);
531                self.bench_results.total_num_views =
532                    metrics.total_num_views.min(cur_metrics.total_num_views);
533                self.bench_results.failed_num_views =
534                    metrics.failed_num_views.max(cur_metrics.failed_num_views);
535            }
536        }
537        self.nodes_post_results += 1;
538        if self.bench_results.partial_results == "Unset" {
539            self.bench_results.partial_results = "One".to_string();
540            self.bench_results.printout();
541            self.output_to_csv();
542        }
543        if self.bench_results.partial_results == "One"
544            && self.nodes_post_results >= (self.config.config.da_staked_committee_size as u64 / 2)
545        {
546            self.bench_results.partial_results = "HalfDA".to_string();
547            self.bench_results.printout();
548            self.output_to_csv();
549        }
550        if self.bench_results.partial_results == "HalfDA"
551            && self.nodes_post_results >= (self.config.config.num_nodes_with_stake.get() as u64 / 2)
552        {
553            self.bench_results.partial_results = "Half".to_string();
554            self.bench_results.printout();
555            self.output_to_csv();
556        }
557        if self.bench_results.partial_results != "Full"
558            && self.nodes_post_results >= (self.config.config.num_nodes_with_stake.get() as u64)
559        {
560            self.bench_results.partial_results = "Full".to_string();
561            self.bench_results.printout();
562            self.output_to_csv();
563        }
564        Ok(())
565    }
566
567    fn post_builder(&mut self, builder: Url) -> Result<(), ServerError> {
568        self.builders.push(builder);
569        Ok(())
570    }
571
572    fn get_builders(&self) -> Result<Vec<Url>, ServerError> {
573        if !matches!(self.config.builder, BuilderType::External)
574            && self.builders.len() != self.config.config.da_staked_committee_size
575        {
576            return Err(ServerError {
577                status: StatusCode::NOT_FOUND,
578                message: "Not all builders are registered yet".to_string(),
579            });
580        }
581        Ok(self.builders.clone())
582    }
583}
584
585/// Shared, lock-guarded orchestrator state, cloned into every axum handler via `State`.
586type SharedOrchestratorState<TYPES> = Arc<RwLock<OrchestratorState<TYPES>>>;
587
588fn malformed_body() -> ServerError {
589    ServerError {
590        status: StatusCode::BAD_REQUEST,
591        message: "Malformed body".to_string(),
592    }
593}
594
595/// Decode a POST body from `OrchestratorClient`.
596///
597/// `client.rs` builds these bodies by calling `body_binary` on an already vbs-serialized
598/// `Vec<u8>`, so the wire body is `[4-byte version][8-byte bincode `Vec<u8>` length][inner
599/// vbs-serialized payload]`; the first 12 bytes are an artifact of that double encoding and must
600/// be stripped before decoding the inner payload.
601fn decode_wrapped_body<T: DeserializeOwned>(body: &[u8]) -> Result<T, ServerError> {
602    body.get(12..)
603        .and_then(|inner| Serializer::<OrchestratorVersion>::deserialize(inner).ok())
604        .ok_or_else(malformed_body)
605}
606
607/// Like [`decode_wrapped_body`], but for the `ready` route, whose inner payload is a raw
608/// [`PeerConfig::to_bytes`] blob rather than a second vbs-serialized value.
609fn decode_wrapped_peer_config<TYPES: NodeType>(
610    body: &[u8],
611) -> Result<PeerConfig<TYPES>, ServerError> {
612    body.get(12..)
613        .and_then(PeerConfig::<TYPES>::from_bytes)
614        .ok_or_else(malformed_body)
615}
616
617async fn healthcheck(headers: HeaderMap) -> Response {
618    healthcheck_response(&headers)
619}
620
621async fn post_identity<TYPES: NodeType>(
622    State(state): State<SharedOrchestratorState<TYPES>>,
623    headers: HeaderMap,
624) -> Response {
625    let result = state.write().await.post_identity();
626    wire::respond(&headers, result)
627}
628
629async fn post_getconfig<TYPES: NodeType>(
630    State(state): State<SharedOrchestratorState<TYPES>>,
631    Path(node_index): Path<u16>,
632    headers: HeaderMap,
633) -> Response {
634    let result = state.write().await.post_getconfig(node_index);
635    wire::respond(&headers, result)
636}
637
638async fn get_tmp_node_index<TYPES: NodeType>(
639    State(state): State<SharedOrchestratorState<TYPES>>,
640    headers: HeaderMap,
641) -> Response {
642    let result = state.write().await.get_tmp_node_index();
643    wire::respond(&headers, result)
644}
645
646async fn post_pubkey<TYPES: NodeType>(
647    State(state): State<SharedOrchestratorState<TYPES>>,
648    Path(is_da): Path<bool>,
649    headers: HeaderMap,
650    body: Bytes,
651) -> Response {
652    let result = match decode_wrapped_body::<Vec<u8>>(&body) {
653        Ok(mut pubkey) => state.write().await.register_public_key(&mut pubkey, is_da),
654        Err(err) => Err(err),
655    };
656    wire::respond(&headers, result)
657}
658
659async fn peer_pubconfig_ready<TYPES: NodeType>(
660    State(state): State<SharedOrchestratorState<TYPES>>,
661    headers: HeaderMap,
662) -> Response {
663    let result = state.read().await.peer_pub_ready();
664    wire::respond(&headers, result)
665}
666
667async fn post_config_after_peer_collected<TYPES: NodeType>(
668    State(state): State<SharedOrchestratorState<TYPES>>,
669    headers: HeaderMap,
670) -> Response {
671    let result = state.write().await.post_config_after_peer_collected();
672    wire::respond(&headers, result)
673}
674
675async fn post_ready<TYPES: NodeType>(
676    State(state): State<SharedOrchestratorState<TYPES>>,
677    headers: HeaderMap,
678    body: Bytes,
679) -> Response {
680    let result = match decode_wrapped_peer_config::<TYPES>(&body) {
681        Ok(peer_config) => state.write().await.post_ready(&peer_config),
682        Err(err) => Err(err),
683    };
684    wire::respond(&headers, result)
685}
686
687async fn post_manual_start<TYPES: NodeType>(
688    State(state): State<SharedOrchestratorState<TYPES>>,
689    headers: HeaderMap,
690    body: Bytes,
691) -> Response {
692    // The raw body is used as-is, with no framing.
693    let result = state.write().await.post_manual_start(body.to_vec());
694    wire::respond(&headers, result)
695}
696
697async fn get_start<TYPES: NodeType>(
698    State(state): State<SharedOrchestratorState<TYPES>>,
699    headers: HeaderMap,
700) -> Response {
701    let result = state.read().await.get_start();
702    wire::respond(&headers, result)
703}
704
705async fn post_results<TYPES: NodeType>(
706    State(state): State<SharedOrchestratorState<TYPES>>,
707    headers: HeaderMap,
708    body: Bytes,
709) -> Response {
710    // Panics on a malformed body: the only production
711    // caller (`client.rs`'s `post_bench_results`) always sends well-formed JSON.
712    let metrics: BenchResults = serde_json::from_slice(&body).unwrap();
713    let result = state.write().await.post_run_results(metrics);
714    wire::respond(&headers, result)
715}
716
717async fn post_builder<TYPES: NodeType>(
718    State(state): State<SharedOrchestratorState<TYPES>>,
719    headers: HeaderMap,
720    body: Bytes,
721) -> Response {
722    let result = match decode_wrapped_body::<Vec<Url>>(&body) {
723        Ok(urls) => {
724            let mut reachable = urls
725                .into_iter()
726                .map(|url| async {
727                    let client: http_client::Client<ClientErr, OrchestratorVersion> =
728                        http_client::Client::builder(url.clone()).build();
729                    client
730                        .connect(Some(Duration::from_secs(2)))
731                        .await
732                        .then_some(url)
733                })
734                .collect::<FuturesUnordered<_>>()
735                .filter_map(futures::future::ready);
736            match reachable.next().await {
737                Some(url) => state.write().await.post_builder(url),
738                None => Err(ServerError {
739                    status: StatusCode::BAD_REQUEST,
740                    message: "No reachable addresses".to_string(),
741                }),
742            }
743        },
744        Err(err) => Err(err),
745    };
746    wire::respond(&headers, result)
747}
748
749async fn get_builders<TYPES: NodeType>(
750    State(state): State<SharedOrchestratorState<TYPES>>,
751    headers: HeaderMap,
752) -> Response {
753    let result = state.read().await.get_builders();
754    wire::respond(&headers, result)
755}
756
757/// Builds the `api` module's routes. Existing clients call these both directly (e.g.
758/// `api/identity`) and under a major-version prefix (`v0/api/identity`), so the `/api` tree is
759/// mounted at the root and under `/v0`.
760fn api_router<TYPES: NodeType>() -> Router<SharedOrchestratorState<TYPES>> {
761    Router::new()
762        .route("/healthcheck", get(healthcheck))
763        .route("/identity", post(post_identity::<TYPES>))
764        .route("/config/{node_index}", post(post_getconfig::<TYPES>))
765        .route("/get_tmp_node_index", post(get_tmp_node_index::<TYPES>))
766        .route("/pubkey/{is_da}", post(post_pubkey::<TYPES>))
767        .route("/peer_pub_ready", get(peer_pubconfig_ready::<TYPES>))
768        .route(
769            "/post_config_after_peer_collected",
770            post(post_config_after_peer_collected::<TYPES>),
771        )
772        .route("/ready", post(post_ready::<TYPES>))
773        .route("/start", get(get_start::<TYPES>))
774        .route("/results", post(post_results::<TYPES>))
775        .route("/manual_start", post(post_manual_start::<TYPES>))
776        .route("/builders", get(get_builders::<TYPES>))
777        .route("/builder", post(post_builder::<TYPES>))
778}
779
780/// Builds the full router: the app-level `healthcheck`, plus the `api` module served both
781/// unversioned and under `/v0`. Every response carries permissive CORS headers.
782fn app<TYPES: NodeType>(state: SharedOrchestratorState<TYPES>) -> Router {
783    let api = Router::new().nest("/api", api_router::<TYPES>());
784    Router::new()
785        .route("/healthcheck", get(healthcheck))
786        .merge(api.clone())
787        .nest("/v0", api)
788        .with_state(state)
789        .layer(cors_layer())
790}
791
792/// Runs the orchestrator
793/// # Errors
794/// This errors if axum runs into an issue during serving
795pub async fn run_orchestrator<TYPES: NodeType>(
796    mut network_config: NetworkConfig<TYPES>,
797    url: Url,
798) -> io::Result<()> {
799    let env_password = std::env::var("ORCHESTRATOR_MANUAL_START_PASSWORD");
800
801    if env_password.is_ok() {
802        tracing::warn!(
803            "Took orchestrator manual start password from the environment variable: \
804             ORCHESTRATOR_MANUAL_START_PASSWORD={:?}",
805            env_password
806        );
807        network_config.manual_start_password = env_password.ok();
808    }
809
810    // Try to overwrite the network_config public keys
811    // from the file the env var points to, or panic.
812    {
813        let env_public_keys = std::env::var("ORCHESTRATOR_PUBLIC_KEYS");
814
815        if let Ok(filepath) = env_public_keys {
816            #[allow(clippy::panic)]
817            let config_file_as_string: String = fs::read_to_string(filepath.clone())
818                .unwrap_or_else(|_| panic!("Could not read config file located at {filepath}"));
819
820            let file: PublicKeysFile<TYPES> =
821                toml::from_str::<PublicKeysFile<TYPES>>(&config_file_as_string)
822                    .expect("Unable to convert config file to TOML");
823
824            network_config.public_keys = file.public_keys;
825        }
826    }
827
828    network_config.config.known_nodes_with_stake = network_config
829        .public_keys
830        .iter()
831        .map(|keys| PeerConfig {
832            stake_table_entry: keys
833                .stake_table_key
834                .stake_table_entry(U256::from(keys.stake)),
835            state_ver_key: keys.state_ver_key.clone(),
836            connect_info: keys.connect_info.clone(),
837        })
838        .collect();
839
840    network_config.config.known_da_nodes = network_config
841        .public_keys
842        .iter()
843        .filter(|keys| keys.da)
844        .map(|keys| PeerConfig {
845            stake_table_entry: keys
846                .stake_table_key
847                .stake_table_entry(U256::from(keys.stake)),
848            state_ver_key: keys.state_ver_key.clone(),
849            connect_info: keys.connect_info.clone(),
850        })
851        .collect();
852
853    let state: SharedOrchestratorState<TYPES> =
854        Arc::new(RwLock::new(OrchestratorState::new(network_config)));
855    let app = app::<TYPES>(state);
856
857    let host = url
858        .host_str()
859        .ok_or_else(|| io::Error::other(format!("orchestrator url missing host: {url}")))?;
860    let port = url
861        .port_or_known_default()
862        .ok_or_else(|| io::Error::other(format!("orchestrator url missing port: {url}")))?;
863    let listener = TcpListener::bind((host, port)).await?;
864
865    tracing::error!("listening on {url:?}");
866    axum::serve(listener, app).await
867}
868
869#[cfg(test)]
870mod tests {
871    use axum::http::{Request, header};
872    use hotshot_example_types::node_types::TestTypes;
873
874    use super::*;
875
876    fn test_app() -> Router {
877        app::<TestTypes>(Arc::new(RwLock::new(OrchestratorState::new(
878            NetworkConfig::default(),
879        ))))
880    }
881
882    /// Every response must carry permissive CORS headers.
883    #[tokio::test]
884    async fn responses_carry_cors_headers() {
885        for uri in ["/healthcheck", "/api/peer_pub_ready", "/no/such/route"] {
886            let req = Request::builder()
887                .uri(uri)
888                .header(header::ORIGIN, "https://example.com")
889                .body(axum::body::Body::empty())
890                .unwrap();
891            let resp = tower::ServiceExt::oneshot(test_app(), req).await.unwrap();
892            assert_eq!(
893                resp.headers()
894                    .get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
895                    .unwrap_or_else(|| panic!("no CORS header on {uri}")),
896                "*",
897                "{uri}"
898            );
899        }
900    }
901
902    /// The `api` module is served both unversioned and under `/v0`; both forms must route.
903    #[tokio::test]
904    async fn versioned_and_unversioned_api_paths_route() {
905        for uri in ["/api/peer_pub_ready", "/v0/api/peer_pub_ready"] {
906            let req = Request::builder()
907                .uri(uri)
908                .body(axum::body::Body::empty())
909                .unwrap();
910            let resp = tower::ServiceExt::oneshot(test_app(), req).await.unwrap();
911            assert_ne!(resp.status(), StatusCode::NOT_FOUND, "{uri} did not route");
912        }
913    }
914}