Skip to content

Commit

Permalink
linera-service: move wallet-only operations out of ClientContext
Browse files Browse the repository at this point in the history
  • Loading branch information
Twey committed Jun 12, 2024
1 parent da21c6a commit f03f018
Show file tree
Hide file tree
Showing 5 changed files with 141 additions and 125 deletions.
54 changes: 38 additions & 16 deletions linera-service/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,15 @@

use std::{
io::{BufRead, BufReader, BufWriter, Write},
iter::IntoIterator,
path::{Path, PathBuf},
};

use anyhow::{bail, Context as _};
use fs4::FileExt as _;
use fs_err::{self, File, OpenOptions};
use linera_base::{
crypto::{BcsSignable, KeyPair, PublicKey},
crypto::{BcsSignable, CryptoRng, KeyPair, PublicKey},
data_types::{Amount, Timestamp},
identifiers::{ChainDescription, ChainId},
};
Expand All @@ -24,7 +25,7 @@ use linera_storage::Storage;
use linera_views::views::ViewError;
use serde::{de::DeserializeOwned, Deserialize, Serialize};

use crate::wallet::Wallet;
use crate::wallet::{UserChain, Wallet};

pub trait Import: DeserializeOwned {
fn read(path: &Path) -> Result<Self, std::io::Error> {
Expand Down Expand Up @@ -125,10 +126,17 @@ impl Drop for FileLock {
/// two processes accessing it at the same time.
pub struct WalletState {
inner: Wallet,
prng: Box<dyn CryptoRng>,
wallet_path: PathBuf,
_lock: FileLock,
}

impl Extend<UserChain> for WalletState {
fn extend<Chains: IntoIterator<Item = UserChain>>(&mut self, chains: Chains) {
self.inner.extend(chains);
}
}

impl WalletState {
pub fn inner(&self) -> &Wallet {
&self.inner
Expand All @@ -145,8 +153,9 @@ impl WalletState {
pub fn from_file(path: &Path) -> Result<Self, anyhow::Error> {
let file = OpenOptions::new().read(true).write(true).open(path)?;
let file_lock = FileLock::new(file, path)?;
let inner = serde_json::from_reader(BufReader::new(&file_lock.file))?;
let inner: Wallet = serde_json::from_reader(BufReader::new(&file_lock.file))?;
Ok(Self {
prng: inner.make_prng(),
inner,
wallet_path: path.into(),
_lock: file_lock,
Expand All @@ -161,20 +170,18 @@ impl WalletState {
let file = Self::open_options().read(true).open(path)?;
let file_lock = FileLock::new(file, path)?;
let mut reader = BufReader::new(&file_lock.file);
if reader.fill_buf()?.is_empty() {
Ok(Self {
inner: Wallet::new(genesis_config, testing_prng_seed),
wallet_path: path.into(),
_lock: file_lock,
})
let inner = if reader.fill_buf()?.is_empty() {
Wallet::new(genesis_config, testing_prng_seed)
} else {
let inner = serde_json::from_reader(reader)?;
Ok(Self {
inner,
wallet_path: path.into(),
_lock: file_lock,
})
}
serde_json::from_reader(reader)?
};

Ok(Self {
prng: inner.make_prng(),
inner,
wallet_path: path.into(),
_lock: file_lock,
})
}

/// Writes the wallet to disk.
Expand Down Expand Up @@ -203,6 +210,21 @@ impl WalletState {
Ok(())
}

pub fn generate_key_pair(&mut self) -> KeyPair {
KeyPair::generate_from(&mut self.prng)
}

pub fn save(&mut self) -> anyhow::Result<()> {
self.inner.refresh_prng_seed(&mut self.prng);
self.write()?;
tracing::info!("Saved user chain states");
Ok(())
}

pub fn refresh_prng_seed(&mut self) {
self.inner.refresh_prng_seed(&mut self.prng)
}

/// Returns options for opening and writing to the wallet file, creating it if it doesn't
/// exist. On Unix, this restricts read and write permissions to the current user.
// TODO(#1924): Implement better key management.
Expand Down
78 changes: 12 additions & 66 deletions linera-service/src/linera/client_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use async_trait::async_trait;
use colored::Colorize;
use futures::{lock::OwnedMutexGuard, Future};
use linera_base::{
crypto::{CryptoRng, KeyPair},
crypto::KeyPair,
data_types::{BlockHeight, HashedBlob, Timestamp},
identifiers::{Account, BlobId, BytecodeId, ChainId},
ownership::ChainOwnership,
Expand All @@ -27,7 +27,7 @@ use linera_execution::Bytecode;
use linera_rpc::node_provider::{NodeOptions, NodeProvider};
use linera_service::{
chain_listener,
config::{GenesisConfig, WalletState},
config::WalletState,
node_service::wait_for_next_round,
storage::StorageConfigNamespace,
wallet::{UserChain, Wallet},
Expand Down Expand Up @@ -65,13 +65,12 @@ use {
use crate::{client_options::ChainOwnershipConfig, ClientOptions};

pub struct ClientContext {
wallet_state: WalletState,
chain_client_builder: ChainClientBuilder<NodeProvider>,
send_timeout: Duration,
recv_timeout: Duration,
notification_retry_delay: Duration,
notification_retries: u32,
prng: Box<dyn CryptoRng>,
pub(crate) wallet_state: WalletState,
pub(crate) chain_client_builder: ChainClientBuilder<NodeProvider>,
pub(crate) send_timeout: Duration,
pub(crate) recv_timeout: Duration,
pub(crate) notification_retry_delay: Duration,
pub(crate) notification_retries: u32,
}

#[async_trait]
Expand Down Expand Up @@ -104,36 +103,8 @@ impl chain_listener::ClientContext<NodeProvider> for ClientContext {
}

impl ClientContext {
pub fn create(
options: &ClientOptions,
genesis_config: GenesisConfig,
testing_prng_seed: Option<u64>,
chains: Vec<UserChain>,
) -> Result<Self, anyhow::Error> {
let wallet_state_path = match &options.wallet_state_path {
Some(path) => path.clone(),
None => Self::create_default_wallet_path()?,
};
anyhow::ensure!(
!wallet_state_path.exists(),
"Wallet already exists at {}. Aborting",
wallet_state_path.display()
);
let mut wallet_state =
WalletState::create(&wallet_state_path, genesis_config, testing_prng_seed)
.with_context(|| format!("Unable to create wallet at {:?}", &wallet_state_path))?;
chains
.into_iter()
.for_each(|chain| wallet_state.inner_mut().insert(chain));
Ok(Self::configure(options, wallet_state))
}

pub fn from_options(options: &ClientOptions) -> Result<Self, anyhow::Error> {
let wallet_state_path = match &options.wallet_state_path {
Some(path) => path.clone(),
None => Self::create_default_wallet_path()?,
};
let wallet_state = WalletState::from_file(&wallet_state_path)?;
let wallet_state = WalletState::from_file(&options.wallet_path()?)?;
Ok(Self::configure(options, wallet_state))
}

Expand All @@ -148,8 +119,6 @@ impl ClientContext {
}

fn configure(options: &ClientOptions, wallet_state: WalletState) -> Self {
let prng = wallet_state.inner().make_prng();

let node_options = NodeOptions {
send_timeout: options.send_timeout,
recv_timeout: options.recv_timeout,
Expand All @@ -167,26 +136,9 @@ impl ClientContext {
recv_timeout: options.recv_timeout,
notification_retry_delay: options.notification_retry_delay,
notification_retries: options.notification_retries,
prng,
}
}

fn create_default_config_path() -> Result<PathBuf, anyhow::Error> {
let mut config_dir = dirs::config_dir()
.context("Default configuration directory not supported. Please specify a path.")?;
config_dir.push("linera");
if !config_dir.exists() {
debug!("{} does not exist, creating", config_dir.display());
fs_err::create_dir(&config_dir)?;
debug!("{} created.", config_dir.display());
}
Ok(config_dir)
}

fn create_default_wallet_path() -> Result<PathBuf, anyhow::Error> {
Ok(Self::create_default_config_path()?.join("wallet.json"))
}

pub fn storage_config(
options: &ClientOptions,
) -> Result<StorageConfigNamespace, anyhow::Error> {
Expand All @@ -195,7 +147,7 @@ impl ClientContext {
#[cfg(feature = "rocksdb")]
None => {
let storage_config = linera_service::storage::StorageConfig::RocksDb {
path: Self::create_default_config_path()?.join("wallet.db"),
path: options.config_path()?.join("wallet.db"),
};
let namespace = "default".to_string();
Ok(StorageConfigNamespace {
Expand Down Expand Up @@ -259,9 +211,7 @@ impl ClientContext {
}

pub fn save_wallet(&mut self) {
self.wallet_state
.inner_mut()
.refresh_prng_seed(&mut self.prng);
self.wallet_state.refresh_prng_seed();
self.wallet_state
.write()
.expect("Unable to write user chains");
Expand Down Expand Up @@ -435,10 +385,6 @@ impl ClientContext {
Ok(blob_id)
}

pub fn generate_key_pair(&mut self) -> KeyPair {
KeyPair::generate_from(&mut self.prng)
}

/// Applies the given function to the chain client.
///
/// Updates the wallet regardless of the outcome. As long as the function returns a round
Expand Down Expand Up @@ -565,7 +511,7 @@ impl ClientContext {
.context("should have default chain")?;
let mut chain_client = self.make_chain_client(storage.clone(), default_chain_id);
while key_pairs.len() < num_chains {
let key_pair = self.generate_key_pair();
let key_pair = self.wallet_state.generate_key_pair();
let public_key = key_pair.public();
let (epoch, committees) = chain_client.epoch_and_committees(default_chain_id).await?;
let epoch = epoch.context("missing epoch on the default chain")?;
Expand Down
52 changes: 45 additions & 7 deletions linera-service/src/linera/client_options.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

use std::{env, iter, num::NonZeroU16, path::PathBuf, time::Duration};

use anyhow::Error;
use chrono::{DateTime, Utc};
use linera_base::{
crypto::PublicKey,
Expand All @@ -18,12 +17,13 @@ use linera_execution::{
};
use linera_service::{
chain_listener::{ChainListenerConfig, ClientContext as _},
config::WalletState,
storage::{full_initialize_storage, run_with_storage},
util,
};
use linera_views::common::CommonStoreConfig;

use crate::{ClientContext, Job};
use crate::{ClientContext, GenesisConfig, Job};

#[derive(clap::Parser)]
#[command(
Expand Down Expand Up @@ -104,7 +104,7 @@ pub struct ClientOptions {
}

impl ClientOptions {
pub fn init() -> Result<Self, anyhow::Error> {
pub fn init() -> anyhow::Result<Self> {
let mut options = <ClientOptions as clap::Parser>::parse();
let suffix = match options.with_wallet {
None => String::new(),
Expand All @@ -121,7 +121,7 @@ impl ClientOptions {
Ok(options)
}

pub async fn run_command_with_storage(self) -> Result<(), Error> {
pub async fn run_command_with_storage(self) -> anyhow::Result<()> {
let context = ClientContext::from_options(&self)?;
let genesis_config = context.wallet().genesis_config().clone();
let wasm_runtime = self.wasm_runtime.with_wasm_default();
Expand All @@ -145,7 +145,7 @@ impl ClientOptions {
Ok(())
}

pub async fn initialize_storage(&self) -> Result<(), Error> {
pub async fn initialize_storage(&self) -> anyhow::Result<()> {
let context = ClientContext::from_options(self)?;
let genesis_config = context.wallet().genesis_config().clone();
let max_concurrent_queries = self.max_concurrent_queries;
Expand All @@ -161,6 +161,44 @@ impl ClientOptions {
full_initialize_storage(full_storage_config, &genesis_config).await?;
Ok(())
}

pub fn wallet(&self) -> anyhow::Result<WalletState> {
WalletState::from_file(&self.wallet_path()?)
}

pub fn wallet_path(&self) -> anyhow::Result<PathBuf> {
self.wallet_state_path
.clone()
.map(Ok)
.unwrap_or_else(|| Ok(self.config_path()?.join("wallet.json")))
}

pub fn config_path(&self) -> anyhow::Result<PathBuf> {
let mut config_dir = dirs::config_dir().ok_or(anyhow::anyhow!(
"Default configuration directory not supported. Please specify a path."
))?;
config_dir.push("linera");
if !config_dir.exists() {
tracing::debug!("{} does not exist, creating", config_dir.display());
fs_err::create_dir(&config_dir)?;
tracing::debug!("{} created.", config_dir.display());
}
Ok(config_dir)
}

pub fn create_wallet(
&self,
genesis_config: GenesisConfig,
testing_prng_seed: Option<u64>,
) -> anyhow::Result<WalletState> {
let wallet_path = self.wallet_path()?;
anyhow::ensure!(
!wallet_path.exists(),
"Wallet already exists at {}. Aborting",
wallet_path.display()
);
WalletState::create(&wallet_path, genesis_config, testing_prng_seed)
}
}

#[derive(Clone, clap::Subcommand)]
Expand Down Expand Up @@ -912,9 +950,9 @@ pub struct ChainOwnershipConfig {
}

impl TryFrom<ChainOwnershipConfig> for ChainOwnership {
type Error = Error;
type Error = anyhow::Error;

fn try_from(config: ChainOwnershipConfig) -> Result<ChainOwnership, Error> {
fn try_from(config: ChainOwnershipConfig) -> anyhow::Result<ChainOwnership> {
let ChainOwnershipConfig {
super_owner_public_keys,
owner_public_keys,
Expand Down
Loading

0 comments on commit f03f018

Please sign in to comment.