refactor(agent): move keypairs into services

Author: Reisen

src/agent.rs

@@ -159,7 +159,7 @@ impl Agent {
 
         // Spawn the remote keypair loader endpoint for both networks
         jhs.append(
-            &mut state::keypairs::spawn(
+            &mut services::keypairs(
                 self.config.primary_network.rpc_url.clone(),
                 self.config
                     .secondary_network
@@ -183,6 +183,7 @@ pub mod config {
         super::{
             metrics,
             pyth,
+            services,
             solana::network,
             state,
         },
@@ -211,7 +212,7 @@ pub mod config {
         #[serde(default)]
         pub metrics_server:        metrics::Config,
         #[serde(default)]
-        pub remote_keypair_loader: state::keypairs::Config,
+        pub remote_keypair_loader: services::keypairs::Config,
     }
 
     impl Config {

src/agent/services.rs

@@ -1,3 +1,7 @@
-mod notifier;
+pub mod keypairs;
+pub mod notifier;
 
-pub use notifier::notifier;
+pub use {
+    keypairs::keypairs,
+    notifier::notifier,
+};

src/agent/services/keypairs.rs

@@ -0,0 +1,229 @@
+//! Keypairs
+//!
+//! The Keypairs Service allows hotloading keys for the running agent.
+
+use {
+    crate::agent::{
+        solana::network::Network,
+        state::keypairs::Keypairs,
+    },
+    anyhow::{
+        Context,
+        Result,
+    },
+    serde::Deserialize,
+    solana_client::nonblocking::rpc_client::RpcClient,
+    solana_sdk::{
+        commitment_config::CommitmentConfig,
+        signature::Keypair,
+        signer::Signer,
+    },
+    std::{
+        net::SocketAddr,
+        sync::Arc,
+    },
+    tokio::task::JoinHandle,
+    warp::{
+        hyper::StatusCode,
+        reject::Rejection,
+        reply::{
+            self,
+            WithStatus,
+        },
+        Filter,
+    },
+};
+
+const DEFAULT_MIN_KEYPAIR_BALANCE_SOL: u64 = 1;
+
+pub fn default_bind_address() -> SocketAddr {
+    "127.0.0.1:9001"
+        .parse()
+        .expect("INTERNAL: Could not build default remote keypair loader bind address")
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(default)]
+pub struct Config {
+    primary_min_keypair_balance_sol:   u64,
+    secondary_min_keypair_balance_sol: u64,
+    bind_address:                      SocketAddr,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            primary_min_keypair_balance_sol:   DEFAULT_MIN_KEYPAIR_BALANCE_SOL,
+            secondary_min_keypair_balance_sol: DEFAULT_MIN_KEYPAIR_BALANCE_SOL,
+            bind_address:                      default_bind_address(),
+        }
+    }
+}
+
+pub async fn keypairs<S>(
+    primary_rpc_url: String,
+    secondary_rpc_url: Option<String>,
+    config: Config,
+    state: Arc<S>,
+) -> Vec<JoinHandle<()>>
+where
+    S: Keypairs,
+    S: Send + Sync + 'static,
+{
+    let ip = config.bind_address.ip();
+
+    if !ip.is_loopback() {
+        tracing::warn!(
+            bind_address = ?config.bind_address,
+            "Remote key loader: bind address is not localhost. Make sure the access on the selected address is secure.",
+        );
+    }
+
+    let primary_upload_route = {
+        let state = state.clone();
+        let rpc_url = primary_rpc_url.clone();
+        let min_balance = config.primary_min_keypair_balance_sol;
+        warp::path!("primary" / "load_keypair")
+            .and(warp::post())
+            .and(warp::body::content_length_limit(1024))
+            .and(warp::body::json())
+            .and(warp::path::end())
+            .and_then(move |kp: Vec<u8>| {
+                let state = state.clone();
+                let rpc_url = rpc_url.clone();
+                async move {
+                    let response = handle_new_keypair(
+                        state,
+                        Network::Primary,
+                        kp,
+                        min_balance,
+                        rpc_url,
+                        "primary",
+                    )
+                    .await;
+                    Result::<WithStatus<_>, Rejection>::Ok(response)
+                }
+            })
+    };
+
+    let secondary_upload_route = warp::path!("secondary" / "load_keypair")
+        .and(warp::post())
+        .and(warp::body::content_length_limit(1024))
+        .and(warp::body::json())
+        .and(warp::path::end())
+        .and_then(move |kp: Vec<u8>| {
+            let state = state.clone();
+            let rpc_url = secondary_rpc_url.clone();
+            async move {
+                if let Some(rpc_url) = rpc_url {
+                    let min_balance = config.secondary_min_keypair_balance_sol;
+                    let response = handle_new_keypair(
+                        state,
+                        Network::Secondary,
+                        kp,
+                        min_balance,
+                        rpc_url,
+                        "secondary",
+                    )
+                    .await;
+                    Result::<WithStatus<_>, Rejection>::Ok(response)
+                } else {
+                    Result::<WithStatus<_>, Rejection>::Ok(reply::with_status(
+                        "Secondary network is not active",
+                        StatusCode::SERVICE_UNAVAILABLE,
+                    ))
+                }
+            }
+        });
+
+    let http_api_jh = {
+        let (_, serve) = warp::serve(primary_upload_route.or(secondary_upload_route))
+            .bind_with_graceful_shutdown(config.bind_address, async {
+                let _ = crate::agent::EXIT.subscribe().changed().await;
+            });
+        tokio::spawn(serve)
+    };
+
+    // WARNING: All jobs spawned here must report their join handles in this vec
+    vec![http_api_jh]
+}
+
+/// Validate and apply a keypair to the specified mut reference,
+/// hiding errors in logs.
+///
+/// Returns the appropriate HTTP response depending on checks success.
+///
+/// NOTE(2023-03-22): Lifetime bounds are currently necessary
+/// because of https://github.com/rust-lang/rust/issues/63033
+async fn handle_new_keypair<'a, 'b: 'a, S>(
+    state: Arc<S>,
+    network: Network,
+    new_keypair_bytes: Vec<u8>,
+    min_keypair_balance_sol: u64,
+    rpc_url: String,
+    network_name: &'b str,
+) -> WithStatus<&'static str>
+where
+    S: Keypairs,
+{
+    let mut upload_ok = true;
+    match Keypair::from_bytes(&new_keypair_bytes) {
+        Ok(kp) => match validate_keypair(&kp, min_keypair_balance_sol, rpc_url.clone()).await {
+            Ok(()) => {
+                Keypairs::update_keypair(&*state, network, kp).await;
+            }
+            Err(e) => {
+                tracing::warn!(
+                    network = network_name,
+                    error = e.to_string(),
+                    "Remote keypair loader: Keypair failed validation",
+                );
+                upload_ok = false;
+            }
+        },
+        Err(e) => {
+            tracing::warn!(
+                network = network_name,
+                error = e.to_string(),
+                "Remote keypair loader: Keypair failed validation",
+            );
+            upload_ok = false;
+        }
+    }
+
+    if upload_ok {
+        reply::with_status("keypair upload OK", StatusCode::OK)
+    } else {
+        reply::with_status(
+            "Could not upload keypair. See logs for details.",
+            StatusCode::BAD_REQUEST,
+        )
+    }
+}
+
+/// Validate keypair balance before using it in transactions.
+pub async fn validate_keypair(
+    kp: &Keypair,
+    min_keypair_balance_sol: u64,
+    rpc_url: String,
+) -> Result<()> {
+    let c = RpcClient::new_with_commitment(rpc_url, CommitmentConfig::confirmed());
+
+    let balance_lamports = c
+        .get_balance(&kp.pubkey())
+        .await
+        .context("Could not check keypair's balance")?;
+
+    let lamports_in_sol = 1_000_000_000;
+
+    if balance_lamports > min_keypair_balance_sol * lamports_in_sol {
+        Ok(())
+    } else {
+        Err(anyhow::anyhow!(format!(
+            "Keypair {} balance of {} SOL below threshold of {} SOL",
+            kp.pubkey(),
+            balance_lamports as f64 / lamports_in_sol as f64,
+            min_keypair_balance_sol
+        )))
+    }
+}