Skip to content

Commit

Permalink
raise need_password_or_derived_key when key is null
Browse files Browse the repository at this point in the history
  • Loading branch information
@tyrone98
tyrone98 committed Aug 2, 2024
1 parent 9681831 commit a2e7592
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 20 deletions.
144 changes: 125 additions & 19 deletions token-core/tcx/src/handler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -593,7 +593,14 @@ pub fn derive_accounts(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let mut account_responses: Vec<AccountResponse> = vec![];

Expand Down Expand Up @@ -645,7 +652,14 @@ pub(crate) fn export_mnemonic(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

tcx_ensure!(
guard.keystore().derivable(),
Expand Down Expand Up @@ -681,7 +695,14 @@ pub(crate) fn export_private_key(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let curve = CurveType::from_str(&param.curve);
let private_key_bytes = guard
Expand Down Expand Up @@ -724,7 +745,13 @@ pub(crate) fn verify_password(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

if keystore.verify_password(&param.key.clone().unwrap().into()) {
if keystore.verify_password(
&param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
) {
let rsp = GeneralResult {
is_success: true,
error: "".to_owned(),
Expand All @@ -744,7 +771,13 @@ pub(crate) fn delete_keystore(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

if keystore.verify_password(&param.key.clone().unwrap().into()) {
if keystore.verify_password(
&param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
) {
delete_keystore_file(&param.id)?;
map.remove(&param.id);

Expand Down Expand Up @@ -804,7 +837,14 @@ pub(crate) fn sign_tx(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

sign_transaction_internal(&param, guard.keystore_mut())
}
Expand All @@ -819,7 +859,14 @@ pub(crate) fn sign_hashes(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let signatures = param
.data_to_sign
Expand Down Expand Up @@ -850,7 +897,14 @@ pub(crate) fn get_public_keys(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let public_keys: Vec<TypedPublicKey> = param
.derivations
Expand Down Expand Up @@ -896,7 +950,14 @@ pub(crate) fn get_extended_public_keys(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let extended_public_keys = param
.derivations
Expand Down Expand Up @@ -926,7 +987,14 @@ pub(crate) fn sign_message(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

sign_message_internal(&param, guard.keystore_mut())
}
Expand All @@ -940,7 +1008,15 @@ pub(crate) fn sign_psbt(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let psbt_input = PsbtInput::decode(
param
.input
Expand Down Expand Up @@ -969,7 +1045,14 @@ pub(crate) fn sign_psbts(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;
let psbt_inputs = PsbtsInput::decode(
param
.input
Expand Down Expand Up @@ -1161,7 +1244,13 @@ pub(crate) fn backup(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let original = keystore.backup(&param.key.clone().unwrap().into())?;
let original = keystore.backup(
&param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;
let fingerprint = match keystore.meta().source {
Source::Mnemonic | Source::NewMnemonic => Some(fingerprint_from_mnemonic(&original)?),
Source::Private | Source::Wif => Some(fingerprint_from_any_format_pk(&original)?),
Expand Down Expand Up @@ -1254,10 +1343,13 @@ pub(crate) fn sign_authentication_message(data: &[u8]) -> Result<Vec<u8>> {
return Err(anyhow::anyhow!("identity_not_found"));
};

let unlocker = identity_ks
.store()
.crypto
.use_key(&param.key.clone().unwrap().into())?;
let unlocker = identity_ks.store().crypto.use_key(
&param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let signature = identity_ks.identity().sign_authentication_message(
param.access_time,
Expand Down Expand Up @@ -1336,7 +1428,14 @@ pub(crate) fn sign_bls_to_execution_change(data: &[u8]) -> Result<Vec<u8>> {
Some(keystore) => Ok(keystore),
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;
let mut guard = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut guard = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;
let result: SignBlsToExecutionChangeResult =
param.sign_bls_to_execution_change(guard.keystore_mut())?;
encode_message(result)
Expand All @@ -1352,7 +1451,14 @@ pub(crate) fn eth_batch_personal_sign(data: &[u8]) -> Result<Vec<u8>> {
_ => Err(anyhow!("{}", "wallet_not_found")),
}?;

let mut keystore = KeystoreGuard::unlock(keystore, param.key.clone().unwrap().into())?;
let mut keystore = KeystoreGuard::unlock(
keystore,
param
.key
.clone()
.expect("need_password_or_derived_key")
.into(),
)?;

let signatures = batch_personal_sign(keystore.keystore_mut(), param.data, &param.path)?;

Expand Down
2 changes: 1 addition & 1 deletion token-core/tcx/src/migration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ pub(crate) fn migrate_keystore(data: &[u8]) -> Result<Vec<u8>> {
let json_str = fs::read_to_string(file_path)?;
let json = serde_json::from_str::<Value>(&json_str)?;

let key = match param.key.clone().unwrap() {
let key = match param.key.clone().expect("need_password_or_derived_key") {
migrate_keystore_param::Key::Password(password) => tcx_crypto::Key::Password(password),
migrate_keystore_param::Key::DerivedKey(derived_key) => {
tcx_crypto::Key::DerivedKey(derived_key)
Expand Down

0 comments on commit a2e7592

Please sign in to comment.