feat: register on proxy (#285)

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(

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)]

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<P: AsRef<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::<Self>("proxy_server")
+            .expect("Failed to generate configuration instance")
+    }
+}

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(())
+}

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<FileBasedKeystore>,
+        keystore: &FileBasedKeystore,
         address_index: usize,
-    ) -> Result<([u8; 32], String), Box<dyn std::error::Error>> {
+    ) -> 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(),
         ))
     }

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 = ""