Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: register on proxy #285

Merged
merged 1 commit into from
Dec 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion atoma-bin/atoma_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -67,6 +71,8 @@ struct Config {
state: AtomaStateManagerConfig,

daemon: AtomaDaemonConfig,

proxy: ProxyConfig,
}

impl Config {
Expand All @@ -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),
}
}
}
Expand Down Expand Up @@ -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(
Expand Down
1 change: 1 addition & 0 deletions atoma-service/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
50 changes: 50 additions & 0 deletions atoma-service/src/proxy/config.rs
Original file line number Diff line number Diff line change
@@ -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")
}
}
47 changes: 47 additions & 0 deletions atoma-service/src/proxy/mod.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
13 changes: 4 additions & 9 deletions atoma-service/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
))
}
Expand Down
8 changes: 8 additions & 0 deletions config.example.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ""