Skip to content

Commit

Permalink
rusk-wallet: Refactor Wallet struct to make impossible states impossible
Browse files Browse the repository at this point in the history
  • Loading branch information
@welf
welf committed Dec 4, 2024
1 parent bb91963 commit bc482df
Show file tree
Hide file tree
Showing 10 changed files with 641 additions and 262 deletions.
6 changes: 3 additions & 3 deletions rusk-wallet/src/bin/command.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,14 @@ use rusk_wallet::gas::{
DEFAULT_PRICE, MIN_PRICE_DEPLOYMENT,
};
use rusk_wallet::{
Address, Error, Profile, Wallet, EPOCH, MAX_CONTRACT_INIT_ARG_SIZE,
MAX_PROFILES,
Address, Error, Profile, Wallet, WalletPath, EPOCH,
MAX_CONTRACT_INIT_ARG_SIZE, MAX_PROFILES,
};
use wallet_core::BalanceInfo;

use crate::io::prompt;
use crate::settings::Settings;
use crate::{WalletFile, WalletPath};
use crate::WalletFile;

/// Commands that can be run against the Dusk wallet
#[allow(clippy::large_enum_variant)]
Expand Down
48 changes: 32 additions & 16 deletions rusk-wallet/src/bin/interactive.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,13 @@ use std::fmt::Display;

use bip39::{Language, Mnemonic, MnemonicType};
use inquire::{InquireError, Select};

use rusk_wallet::currency::Dusk;
use rusk_wallet::dat::{DatFileVersion, LATEST_VERSION};
use rusk_wallet::{Address, Error, Profile, Wallet, WalletPath, MAX_PROFILES};
use rusk_wallet::{
Address, Error, Profile, SecureWalletFile, Wallet, WalletFilePath,
WalletPath, MAX_PROFILES,
};

use crate::io::{self, prompt};
use crate::settings::Settings;
Expand Down Expand Up @@ -138,6 +142,9 @@ async fn profile_idx(
match menu_profile(wallet)? {
ProfileSelect::Index(index, _) => Ok(index),
ProfileSelect::New => {
// get the wallet file
let file = wallet.file().clone().ok_or(Error::WalletFileMissing)?;

if wallet.profiles().len() >= MAX_PROFILES {
println!(
"Cannot create more profiles, this wallet only supports up to {MAX_PROFILES} profiles"
Expand All @@ -147,24 +154,22 @@ async fn profile_idx(
}

let profile_idx = wallet.add_profile();
let file_version = wallet.get_file_version()?;

let password = &settings.password;
// if the version file is old, ask for password and save as
// latest dat file
if file_version.is_old() {
if file.is_old() {
let pwd = prompt::request_auth(
"Updating your wallet data file, please enter your wallet password ",
password,
DatFileVersion::RuskBinaryFileFormat(LATEST_VERSION),
)?;

// UNWRAP: we can safely unwrap here because we know the file is
// not None since we've checked the file version
wallet.save_to(WalletFile {
path: wallet.file().clone().unwrap().path,
wallet.save_to(WalletFile::new(
file.path().clone(),
pwd,
})?;
DatFileVersion::RuskBinaryFileFormat(LATEST_VERSION),
))?;
} else {
// else just save
wallet.save()?;
Expand Down Expand Up @@ -231,8 +236,10 @@ pub(crate) async fn load_wallet(
settings: &Settings,
file_version: Result<DatFileVersion, Error>,
) -> anyhow::Result<Wallet<WalletFile>> {
let wallet_found =
wallet_path.inner().exists().then(|| wallet_path.clone());
let wallet_found = wallet_path
.wallet_path()
.exists()
.then(|| wallet_path.clone());

let password = &settings.password;

Expand All @@ -247,10 +254,11 @@ pub(crate) async fn load_wallet(
password,
file_version,
)?;
match Wallet::from_file(WalletFile {
path: path.clone(),
match Wallet::from_file(WalletFile::new(
path.clone(),
pwd,
}) {
file_version,
)) {
Ok(wallet) => break wallet,
Err(_) if attempt > 2 => {
Err(Error::AttemptsExhausted)?;
Expand All @@ -277,7 +285,11 @@ pub(crate) async fn load_wallet(
// create and store the wallet
let mut w = Wallet::new(mnemonic)?;
let path = wallet_path.clone();
w.save_to(WalletFile { path, pwd })?;
w.save_to(WalletFile::new(
path,
pwd,
DatFileVersion::RuskBinaryFileFormat(LATEST_VERSION),
))?;
w
}
MainMenu::Recover => {
Expand All @@ -292,7 +304,11 @@ pub(crate) async fn load_wallet(
// create and store the recovered wallet
let mut w = Wallet::new(phrase)?;
let path = wallet_path.clone();
w.save_to(WalletFile { path, pwd })?;
w.save_to(WalletFile::new(
path,
pwd,
DatFileVersion::RuskBinaryFileFormat(LATEST_VERSION),
))?;
w
}
MainMenu::Exit => std::process::exit(0),
Expand Down Expand Up @@ -493,7 +509,7 @@ impl Display for MainMenu {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MainMenu::Load(path) => {
write!(f, "Load wallet from {}", path.wallet.display())
write!(f, "Load wallet from {}", path.wallet_path().display())
}
MainMenu::Create => write!(f, "Create a new wallet"),
MainMenu::Recover => {
Expand Down
57 changes: 20 additions & 37 deletions rusk-wallet/src/bin/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ mod io;
mod settings;

pub(crate) use command::{Command, RunResult};
use rusk_wallet::{WalletFilePath, WalletPath};

use std::fs::{self, File};
use std::io::Write;
Expand All @@ -19,35 +20,20 @@ use bip39::{Language, Mnemonic, MnemonicType};
use clap::Parser;
use inquire::InquireError;
use rocksdb::ErrorKind;
use rusk_wallet::currency::Dusk;
use rusk_wallet::dat::{self, LATEST_VERSION};
use rusk_wallet::{
Error, GraphQL, Profile, SecureWalletFile, Wallet, WalletPath, EPOCH,
};
use tracing::{error, info, warn, Level};

use crate::command::TransactionHistory;
use crate::settings::{LogFormat, Settings};

use rusk_wallet::{
currency::Dusk,
dat::{self, LATEST_VERSION},
Error, GraphQL, Profile, SecureWalletFile, Wallet, WalletFile, EPOCH,
};

use config::Config;
use io::{prompt, status, WalletArgs};

#[derive(Debug, Clone)]
pub(crate) struct WalletFile {
path: WalletPath,
pwd: Vec<u8>,
}

impl SecureWalletFile for WalletFile {
fn path(&self) -> &WalletPath {
&self.path
}

fn pwd(&self) -> &[u8] {
&self.pwd
}
}

#[tokio::main(flavor = "multi_thread")]
async fn main() -> anyhow::Result<()> {
if let Err(err) = exec().await {
Expand Down Expand Up @@ -139,7 +125,7 @@ async fn exec() -> anyhow::Result<()> {

// prepare wallet path
let mut wallet_path =
WalletPath::from(wallet_dir.as_path().join("wallet.dat"));
WalletPath::try_from(wallet_dir.as_path().join("wallet.dat"))?;

// load configuration (or use default)
let cfg = Config::load(&wallet_dir)?;
Expand Down Expand Up @@ -193,6 +179,7 @@ async fn exec() -> anyhow::Result<()> {
return Ok(());
};

// get the wallet file version
let file_version = dat::read_file_version(&wallet_path);

// get our wallet ready
Expand Down Expand Up @@ -232,10 +219,7 @@ async fn exec() -> anyhow::Result<()> {
// create wallet
let mut w = Wallet::new(mnemonic)?;

w.save_to(WalletFile {
path: wallet_path,
pwd,
})?;
w.save_to(WalletFile::new(wallet_path, pwd, file_version?))?;

w
}
Expand All @@ -253,10 +237,11 @@ async fn exec() -> anyhow::Result<()> {
file_version,
)?;

let w = Wallet::from_file(WalletFile {
path: file.clone(),
pwd: pwd.clone(),
})?;
let w = Wallet::from_file(WalletFile::new(
file.clone(),
pwd.clone(),
file_version,
))?;

(w, pwd)
}
Expand All @@ -279,10 +264,7 @@ async fn exec() -> anyhow::Result<()> {
}
};

w.save_to(WalletFile {
path: wallet_path,
pwd,
})?;
w.save_to(WalletFile::new(wallet_path, pwd, file_version?))?;

w
}
Expand All @@ -297,10 +279,11 @@ async fn exec() -> anyhow::Result<()> {
file_version,
)?;

Wallet::from_file(WalletFile {
path: wallet_path,
Wallet::from_file(WalletFile::new(
wallet_path,
pwd,
})?
file_version,
))?
}
},
};
Expand Down
94 changes: 5 additions & 89 deletions rusk-wallet/src/dat.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,7 @@
use std::fs;
use std::io::Read;

use wallet_core::Seed;

use crate::crypto::decrypt;
use crate::{Error, WalletPath};
use crate::{Error, WalletFilePath, WalletPath};

/// Binary prefix for old Dusk wallet files
pub const OLD_MAGIC: u32 = 0x1d0c15;
Expand All @@ -38,89 +35,6 @@ pub enum DatFileVersion {
RuskBinaryFileFormat(Version),
}

impl DatFileVersion {
/// Checks if the file version is older than the latest Rust Binary file
/// format
pub fn is_old(&self) -> bool {
matches!(self, Self::Legacy | Self::OldWalletCli(_))
}
}

/// Make sense of the payload and return it
pub(crate) fn get_seed_and_address(
file: DatFileVersion,
mut bytes: Vec<u8>,
pwd: &[u8],
) -> Result<(Seed, u8), Error> {
match file {
DatFileVersion::Legacy => {
if bytes[1] == 0 && bytes[2] == 0 {
bytes.drain(..3);
}

bytes = decrypt(&bytes, pwd)?;

// get our seed
let seed = bytes[..]
.try_into()
.map_err(|_| Error::WalletFileCorrupted)?;

Ok((seed, 1))
}
DatFileVersion::OldWalletCli((major, minor, _, _, _)) => {
bytes.drain(..5);

let result: Result<(Seed, u8), Error> = match (major, minor) {
(1, 0) => {
let content = decrypt(&bytes, pwd)?;
let buff = &content[..];

let seed = buff
.try_into()
.map_err(|_| Error::WalletFileCorrupted)?;

Ok((seed, 1))
}
(2, 0) => {
let content = decrypt(&bytes, pwd)?;
let buff = &content[..];

// extract seed
let seed = buff
.try_into()
.map_err(|_| Error::WalletFileCorrupted)?;

// extract addresses count
Ok((seed, buff[0]))
}
_ => Err(Error::UnknownFileVersion(major, minor)),
};

result
}
DatFileVersion::RuskBinaryFileFormat(_) => {
let rest = bytes.get(12..(12 + 96));
if let Some(rest) = rest {
let content = decrypt(rest, pwd)?;

if let Some(seed_buff) = content.get(0..65) {
let seed = seed_buff[0..64]
.try_into()
.map_err(|_| Error::WalletFileCorrupted)?;

let count = &seed_buff[64..65];

Ok((seed, count[0]))
} else {
Err(Error::WalletFileCorrupted)
}
} else {
Err(Error::WalletFileCorrupted)
}
}
}
}

/// From the first 12 bytes of the file (header), we check version
///
/// https://github.com/dusk-network/rusk/wiki/Binary-File-Format/#header
Expand Down Expand Up @@ -193,8 +107,10 @@ pub(crate) fn check_version(

/// Read the first 12 bytes of the dat file and get the file version from
/// there
pub fn read_file_version(file: &WalletPath) -> Result<DatFileVersion, Error> {
let path = &file.wallet;
pub fn read_file_version(
wallet_file_path: &WalletPath,
) -> Result<DatFileVersion, Error> {
let path = &wallet_file_path.wallet_path();

// make sure file exists
if !path.is_file() {
Expand Down
3 changes: 3 additions & 0 deletions rusk-wallet/src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,9 @@ pub enum Error {
/// Contract file location not found
#[error("Invalid WASM contract path provided")]
InvalidWasmContractPath,
/// Invalid wallet file path
#[error("Invalid wallet file path")]
InvalidWalletFilePath,
/// Invalid environment variable value
#[error("Invalid environment variable value {0}")]
InvalidEnvVar(String),
Expand Down
3 changes: 2 additions & 1 deletion rusk-wallet/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,8 @@ pub use error::Error;
pub use gql::{BlockTransaction, GraphQL};
pub use rues::RuesHttpClient;
pub use wallet::{
Address, DecodedNote, Profile, SecureWalletFile, Wallet, WalletPath,
Address, DecodedNote, Profile, SecureWalletFile, Wallet, WalletFile,
WalletFilePath, WalletPath,
};

use execution_core::stake::StakeData;
Expand Down
Loading

0 comments on commit bc482df

Please sign in to comment.