Merge pull request #90 from monadicus/feat-agent-privatekey

feat(agent): local private keys

crates/aot/src/runner/mod.rs

@@ -4,8 +4,9 @@ use std::{
 };
 
 use aleo_std::StorageMode;
-use anyhow::Result;
+use anyhow::{bail, Result};
 use clap::Args;
+use core::str::FromStr;
 use serde::{Deserialize, Serialize};
 use serde_clap_deserialize::serde_clap_default;
 use snarkos_node::Node;
@@ -16,6 +17,31 @@ use crate::{ledger::Addrs, Account, PrivateKey};
 
 mod metrics;
 
+#[derive(Debug, Args, Serialize, Deserialize)]
+#[group(required = true, multiple = false)]
+pub struct Key {
+    /// Specify the account private key of the node
+    #[clap(long = "private-key")]
+    pub private_key: Option<PrivateKey>,
+    /// Specify the account private key of the node
+    #[clap(long = "private-key-file")]
+    pub private_key_file: Option<PathBuf>,
+}
+
+impl Key {
+    pub fn try_get(self) -> Result<PrivateKey> {
+        match (self.private_key, self.private_key_file) {
+            (Some(key), None) => Ok(key),
+            (None, Some(file)) => {
+                let raw = std::fs::read_to_string(file)?.trim().to_string();
+                Ok(PrivateKey::from_str(&raw)?)
+            }
+            // clap should make this unreachable, but serde might not
+            _ => bail!("Either `private-key` or `private-key-file` must be set"),
+        }
+    }
+}
+
 #[serde_clap_default]
 #[derive(Debug, Args, Serialize, Deserialize)]
 pub struct Runner {
@@ -30,9 +56,8 @@ pub struct Runner {
     #[arg(required = true, name = "type", short, long)]
     pub node_type: NodeType,
 
-    /// Specify the account private key of the node
-    #[clap(long = "private-key")]
-    pub private_key: PrivateKey,
+    #[clap(flatten)]
+    pub key: Key,
 
     #[clap(long = "bind", default_value_t = IpAddr::V4(Ipv4Addr::UNSPECIFIED))]
     pub bind_addr: IpAddr,
@@ -70,7 +95,7 @@ impl Runner {
         let bft_ip = SocketAddr::new(bind_addr, self.bft);
         let metrics_ip = SocketAddr::new(bind_addr, self.metrics);
 
-        let account = Account::try_from(self.private_key)?;
+        let account = Account::try_from(self.key.try_get()?)?;
 
         let genesis = Block::from_bytes_le(&std::fs::read(&self.genesis)?)?;
         let storage_mode = StorageMode::Custom(self.ledger);

crates/snot-agent/src/cli.rs

@@ -1,13 +1,13 @@
 use std::{
-    env,
+    env, fs,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
 };
 
 use clap::Parser;
 use http::Uri;
 use snot_common::state::{AgentId, AgentMode, PortConfig};
-use tracing::info;
+use tracing::{info, warn};
 
 pub const ENV_ENDPOINT: &str = "SNOT_ENDPOINT";
 pub const ENV_ENDPOINT_DEFAULT: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1234);
@@ -23,6 +23,11 @@ pub struct Cli {
     #[arg(long)]
     pub id: Option<AgentId>,
 
+    /// Locally provided private key file, used for envs where private keys are locally provided
+    #[arg(long)]
+    #[clap(long = "private-key-file")]
+    pub private_key_file: Option<PathBuf>,
+
     #[arg(long, value_delimiter = ',', num_args = 1..)]
     pub labels: Option<Vec<String>>,
 
@@ -60,12 +65,21 @@ impl Cli {
 
         let mut query = format!("/agent?mode={}", u8::from(self.modes));
 
-        // add ?id=
+        // add &id=
         if let Some(id) = self.id {
             query.push_str(&format!("&id={}", id));
         }
 
-        // add ?labels= or &labels= if id is present
+        // add local pk flag
+        if let Some(file) = self.private_key_file.as_ref() {
+            if fs::metadata(file).is_ok() {
+                query.push_str("&local_pk=true");
+            } else {
+                warn!("private-key-file flag ignored as the file was not found: {file:?}")
+            }
+        }
+
+        // add &labels= if id is present
         if let Some(labels) = &self.labels {
             info!("using labels: {:?}", labels);
             query.push_str(&format!("&labels={}", labels.join(",")));

crates/snot-agent/src/rpc.rs

@@ -9,7 +9,7 @@ use snot_common::{
         control::{ControlServiceRequest, ControlServiceResponse},
         MuxMessage,
     },
-    state::{AgentId, AgentPeer, AgentState, PortConfig},
+    state::{AgentId, AgentPeer, AgentState, KeyState, PortConfig},
 };
 use tarpc::{context, ClientMessage, Response};
 use tokio::{
@@ -247,8 +247,20 @@ impl AgentService for AgentRpcServer {
                         .arg("--node")
                         .arg(state.cli.ports.node.to_string());
 
-                    if let Some(pk) = node.private_key {
-                        command.arg("--private-key").arg(pk);
+                    match node.private_key {
+                        KeyState::None => {}
+                        KeyState::Local => {
+                            command.arg("--private-key-file").arg(
+                                state
+                                    .cli
+                                    .private_key_file
+                                    .as_ref()
+                                    .ok_or(ReconcileError::NoLocalPrivateKey)?,
+                            );
+                        }
+                        KeyState::Literal(pk) => {
+                            command.arg("--private-key").arg(pk);
+                        }
                     }
 
                     // Find agents that do not have cached addresses

crates/snot-common/src/lib.rs

@@ -1,9 +1,11 @@
 pub mod rpc;
+pub mod set;
 pub mod state;
 pub use lasso;
 
 pub mod prelude {
     pub use crate::rpc::*;
+    pub use crate::set::*;
     pub use crate::state::*;
 }
 

crates/snot-common/src/rpc/agent.rs

@@ -45,6 +45,8 @@ pub enum ReconcileError {
     StorageAcquireError,
     #[error("failed to resolve addresses of stated peers")]
     ResolveAddrError(ResolveError),
+    #[error("agent did not provide a local private key")]
+    NoLocalPrivateKey,
     #[error("unknown error")]
     Unknown,
 }

crates/snot-common/src/set.rs

@@ -0,0 +1,10 @@
+pub const MASK_PREFIX_LEN: usize = 5;
+
+#[repr(usize)]
+pub enum MaskBit {
+    Validator = 0,
+    Prover = 1,
+    Client = 2,
+    Compute = 3,
+    LocalPrivateKey = 4,
+}

crates/snot-common/src/state.rs

@@ -11,7 +11,7 @@ use lazy_static::lazy_static;
 use regex::Regex;
 use serde::{de::Error, Deserialize, Serialize};
 
-use crate::INTERN;
+use crate::{prelude::MaskBit, INTERN};
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct AgentId(Spur);
@@ -43,14 +43,45 @@ impl AgentState {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct NodeState {
     pub ty: NodeType,
-    pub private_key: Option<String>,
+    pub private_key: KeyState,
     pub height: (usize, HeightRequest),
 
     pub online: bool,
     pub peers: Vec<AgentPeer>,
     pub validators: Vec<AgentPeer>,
 }
 
+/// A representation of which key to use for the agent.
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub enum KeyState {
+    /// No private key provided
+    #[default]
+    None,
+    /// A private key is provided by the agent
+    Local,
+    /// A literal private key
+    Literal(String),
+    // TODO: generated?/new
+}
+
+impl From<Option<String>> for KeyState {
+    fn from(s: Option<String>) -> Self {
+        match s {
+            Some(s) => Self::Literal(s),
+            None => Self::None,
+        }
+    }
+}
+
+impl KeyState {
+    pub fn try_string(&self) -> Option<String> {
+        match self {
+            Self::Literal(s) => Some(s.to_owned()),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, Parser)]
 pub struct PortConfig {
     /// Specify the IP address and port for the node server
@@ -208,11 +239,11 @@ impl NodeType {
     }
 
     pub fn bit(self) -> usize {
-        match self {
-            Self::Validator => 0,
-            Self::Prover => 1,
-            Self::Client => 2,
-        }
+        (match self {
+            Self::Validator => MaskBit::Validator,
+            Self::Prover => MaskBit::Prover,
+            Self::Client => MaskBit::Client,
+        }) as usize
     }
 }
 

crates/snot/src/cannon/source.rs

@@ -197,14 +197,14 @@ impl TxSource {
                 let sample_pk = || {
                     private_keys
                         .get(rand::random::<usize>() % private_keys.len())
-                        .and_then(|k| env.storage.sample_keysource_pk(k))
+                        .and_then(|k| env.storage.sample_keysource_pk(k).try_string())
                         .ok_or(anyhow!("error selecting a valid private key"))
                 };
                 let sample_addr = || {
                     addresses
                         .get(rand::random::<usize>() % addresses.len())
-                        .and_then(|k| env.storage.sample_keysource_addr(k))
-                        .ok_or(anyhow!("error selecting a valid private key"))
+                        .and_then(|k| env.storage.sample_keysource_addr(k).try_string())
+                        .ok_or(anyhow!("error selecting a valid address"))
                 };
 
                 let Some(mode) = tx_modes

crates/snot/src/env/mod.rs

@@ -442,7 +442,8 @@ pub async fn initial_reconcile(env_id: usize, state: &GlobalState) -> anyhow::Re
             node_state.private_key = node
                 .key
                 .as_ref()
-                .and_then(|key| env.storage.lookup_keysource_pk(key));
+                .map(|key| env.storage.lookup_keysource_pk(key))
+                .unwrap_or_default();
 
             let not_me = |agent: &AgentPeer| !matches!(agent, AgentPeer::Internal(candidate_id, _) if *candidate_id == id);
 

crates/snot/src/env/set.rs

@@ -8,6 +8,7 @@ use indexmap::IndexMap;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use snot_common::{
     lasso::Spur,
+    set::MASK_PREFIX_LEN,
     state::{AgentId, NodeKey},
 };
 use thiserror::Error;
@@ -141,8 +142,8 @@ fn _find_compute_agent_by_mask<'a, I: Iterator<Item = &'a Agent>>(
     labels: &[Spur],
 ) -> Option<(&'a Agent, Arc<Busy>)> {
     // replace with
-    let mut mask = FixedBitSet::with_capacity(labels.len() + 4);
-    mask.insert_range(4..labels.len() + 4);
+    let mut mask = FixedBitSet::with_capacity(labels.len() + MASK_PREFIX_LEN);
+    mask.insert_range(MASK_PREFIX_LEN..labels.len() + MASK_PREFIX_LEN);
 
     agents.find_map(|agent| {
         AgentMapping::new(BusyMode::Compute, agent, labels)

crates/snot/src/schema/nodes.rs

@@ -6,7 +6,8 @@ use lazy_static::lazy_static;
 use serde::{de::Visitor, Deserialize, Deserializer, Serialize};
 use snot_common::{
     lasso::Spur,
-    state::{AgentId, HeightRequest, NodeState, NodeType},
+    set::{MaskBit, MASK_PREFIX_LEN},
+    state::{AgentId, HeightRequest, KeyState, NodeState, NodeType},
     INTERN,
 };
 
@@ -91,7 +92,7 @@ impl Node {
     pub fn into_state(&self, ty: NodeType) -> NodeState {
         NodeState {
             ty,
-            private_key: None,
+            private_key: KeyState::None,
 
             // TODO
             height: (0, HeightRequest::Top),
@@ -105,12 +106,20 @@ impl Node {
     }
 
     pub fn mask(&self, key: &NodeKey, labels: &[Spur]) -> FixedBitSet {
-        let mut mask = FixedBitSet::with_capacity(labels.len() + 4);
+        let mut mask = FixedBitSet::with_capacity(labels.len() + MASK_PREFIX_LEN);
+
+        // validator/prover/client
         mask.insert(key.ty.bit());
 
+        // local private key
+        if matches!(self.key, Some(KeySource::Local)) {
+            mask.insert(MaskBit::LocalPrivateKey as usize);
+        }
+
+        // labels
         for (i, label) in labels.iter().enumerate() {
             if self.labels.contains(label) {
-                mask.insert(i + 4);
+                mask.insert(i + MASK_PREFIX_LEN);
             }
         }
         mask
@@ -119,6 +128,8 @@ impl Node {
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum KeySource {
+    /// Private key owned by the agent
+    Local,
     /// APrivateKey1zkp...
     Literal(String),
     /// committee.0 or committee.$ (for replicas)
@@ -167,8 +178,11 @@ impl FromStr for KeySource {
         // use KeySource::Commitee(Option<usize>) when the string is "committee.0" or "committee.$"
         // use KeySource::Named(String, Option<usize>) when the string is "\w+.0" or "\w+.$"
 
+        if s == "local" {
+            return Ok(KeySource::Local);
+        }
         // aleo private key
-        if s.len() == 59 && s.starts_with("APrivateKey1") {
+        else if s.len() == 59 && s.starts_with("APrivateKey1") {
             return Ok(KeySource::Literal(s.to_string()));
 
         // committee key
@@ -208,6 +222,7 @@ impl Display for KeySource {
             f,
             "{}",
             match self {
+                KeySource::Local => "local".to_owned(),
                 KeySource::Literal(key) => key.to_owned(),
                 KeySource::Committee(None) => "committee.$".to_owned(),
                 KeySource::Committee(Some(idx)) => {