From 4ce216aa2bce597faec134e2b57fbddcad830b97 Mon Sep 17 00:00:00 2001 From: Martin Stefcek <35243812+Cifko@users.noreply.github.com> Date: Mon, 16 Dec 2024 10:09:25 +0100 Subject: [PATCH] feat: register on proxy (#285) --- atoma-bin/atoma_node.rs | 14 ++++++++- atoma-service/src/lib.rs | 1 + atoma-service/src/proxy/config.rs | 50 +++++++++++++++++++++++++++++++ atoma-service/src/proxy/mod.rs | 47 +++++++++++++++++++++++++++++ atoma-service/src/server.rs | 13 +++----- config.example.toml | 8 +++++ 6 files changed, 123 insertions(+), 10 deletions(-) create mode 100644 atoma-service/src/proxy/config.rs create mode 100644 atoma-service/src/proxy/mod.rs diff --git a/atoma-bin/atoma_node.rs b/atoma-bin/atoma_node.rs index 6fe8a825..e19e5f91 100644 --- a/atoma-bin/atoma_node.rs +++ b/atoma-bin/atoma_node.rs @@ -3,7 +3,11 @@ use std::{path::Path, str::FromStr, sync::Arc}; use anyhow::{Context, Result}; use atoma_confidential::AtomaConfidentialComputeService; use atoma_daemon::{AtomaDaemonConfig, DaemonState}; -use atoma_service::{config::AtomaServiceConfig, server::AppState}; +use atoma_service::{ + config::AtomaServiceConfig, + proxy::{config::ProxyConfig, register_on_proxy}, + server::AppState, +}; use atoma_state::{config::AtomaStateManagerConfig, AtomaState, AtomaStateManager}; use atoma_sui::{client::AtomaSuiClient, AtomaSuiConfig, SuiEventSubscriber}; use atoma_utils::spawn_with_shutdown; @@ -67,6 +71,8 @@ struct Config { state: AtomaStateManagerConfig, daemon: AtomaDaemonConfig, + + proxy: ProxyConfig, } impl Config { @@ -76,6 +82,7 @@ impl Config { service: AtomaServiceConfig::from_file_path(path), state: AtomaStateManagerConfig::from_file_path(path), daemon: AtomaDaemonConfig::from_file_path(path), + proxy: ProxyConfig::from_file_path(path), } } } @@ -213,6 +220,11 @@ async fn main() -> Result<()> { let client = Arc::new(RwLock::new( AtomaSuiClient::new_from_config(args.config_path).await?, )); + + for (_, node_small_id) in config.daemon.node_badges.iter() { + register_on_proxy(&config.proxy, *node_small_id, &keystore, args.address_index).await?; + } + let (compute_shared_secret_sender, compute_shared_secret_receiver) = tokio::sync::mpsc::unbounded_channel(); let confidential_compute_service = AtomaConfidentialComputeService::new( diff --git a/atoma-service/src/lib.rs b/atoma-service/src/lib.rs index 7f966c6d..26d1768e 100644 --- a/atoma-service/src/lib.rs +++ b/atoma-service/src/lib.rs @@ -5,6 +5,7 @@ pub(crate) mod components; pub mod config; pub(crate) mod handlers; pub mod middleware; +pub mod proxy; pub mod server; pub mod streamer; #[cfg(test)] diff --git a/atoma-service/src/proxy/config.rs b/atoma-service/src/proxy/config.rs new file mode 100644 index 00000000..ab368203 --- /dev/null +++ b/atoma-service/src/proxy/config.rs @@ -0,0 +1,50 @@ +use config::{Config, File}; +use serde::Deserialize; +use std::path::Path; + +/// Configuration for the proxy server +/// +/// This struct holds the configuration parameters needed to connect to proxy server. +#[derive(Debug, Deserialize)] +pub struct ProxyConfig { + pub proxy_address: String, + pub node_public_address: String, + pub country: String, +} + +impl ProxyConfig { + /// Creates a new ProxyConfig instance from a configuration file + /// + /// # Arguments + /// + /// * `config_file_path` - Path to the configuration file. The file should be in a format + /// supported by the `config` crate (e.g., TOML, JSON, YAML) and + /// contain an "proxy_server" section with the required configuration + /// parameters. + /// + /// # Returns + /// + /// Returns a new `ProxyConfig` instance populated with values from the config file. + /// + /// # Panics + /// + /// This method will panic if: + /// * The configuration file cannot be read or parsed + /// * The "proxy_server" section is missing from the configuration + /// * The configuration format doesn't match the expected structure + pub fn from_file_path>(config_file_path: P) -> Self { + let builder = Config::builder() + .add_source(File::with_name(config_file_path.as_ref().to_str().unwrap())) + .add_source( + config::Environment::with_prefix("PROXY_SERVER") + .keep_prefix(true) + .separator("__"), + ); + let config = builder + .build() + .expect("Failed to generate atoma-daemon configuration file"); + config + .get::("proxy_server") + .expect("Failed to generate configuration instance") + } +} diff --git a/atoma-service/src/proxy/mod.rs b/atoma-service/src/proxy/mod.rs new file mode 100644 index 00000000..f1e1c69f --- /dev/null +++ b/atoma-service/src/proxy/mod.rs @@ -0,0 +1,47 @@ +use atoma_utils::constants::SIGNATURE; +use config::ProxyConfig; +use reqwest::Client; +use serde_json::json; +use sui_keys::keystore::FileBasedKeystore; + +use crate::server::utils::sign_response_body; + +pub mod config; + +/// Registers the node on the proxy server +/// +/// # Arguments +/// +/// * `config` - Proxy configuration +/// * `node_small_id` - Small ID of the node +/// * `keystore` - Keystore for signing the registration request +/// * `address_index` - Index of the address to use for signing +pub async fn register_on_proxy( + config: &ProxyConfig, + node_small_id: u64, + keystore: &FileBasedKeystore, + address_index: usize, +) -> anyhow::Result<()> { + let client = Client::new(); + let url = format!("{}/node/registration", config.proxy_address); + + let body = json!({ + "node_small_id": node_small_id, + "public_address": config.node_public_address, + "country": config.country, + }); + + let (_, signature) = sign_response_body(&body, keystore, address_index)?; + + let res = client + .post(&url) + .header(SIGNATURE, signature) + .json(&body) + .send() + .await?; + + if !res.status().is_success() { + anyhow::bail!("Failed to register on proxy server: {}", res.status()); + } + Ok(()) +} diff --git a/atoma-service/src/server.rs b/atoma-service/src/server.rs index 02572bef..78445bee 100644 --- a/atoma-service/src/server.rs +++ b/atoma-service/src/server.rs @@ -395,21 +395,16 @@ pub(crate) mod utils { /// * The SHA-256 hash cannot be converted to a 32-byte array pub(crate) fn sign_response_body( response_body: &Value, - keystore: &Arc, + keystore: &FileBasedKeystore, address_index: usize, - ) -> Result<([u8; 32], String), Box> { + ) -> anyhow::Result<([u8; 32], String)> { let address = keystore.addresses()[address_index]; let response_body_str = response_body.to_string(); let response_body_bytes = response_body_str.as_bytes(); let blake2b_hash = blake2b_hash(response_body_bytes); - let signature = keystore - .sign_hashed(&address, blake2b_hash.as_slice()) - .expect("Failed to sign response body"); + let signature = keystore.sign_hashed(&address, blake2b_hash.as_slice())?; Ok(( - blake2b_hash - .as_slice() - .try_into() - .expect("Invalid BLAKE2b hash length"), + blake2b_hash.as_slice().try_into()?, signature.encode_base64(), )) } diff --git a/config.example.toml b/config.example.toml index c5862477..2b9d3842 100644 --- a/config.example.toml +++ b/config.example.toml @@ -33,3 +33,11 @@ node_badges = [ 1, ], ] # List of node badges, where each badge is a tuple of (badge_id, small_id), both values are assigned once the node registers itself + +[proxy_server] +# replace this with the address of the proxy server +proxy_address = "" +# replace this with the public address of this node +node_public_address = "" +# replace this with the country of the node +country = ""