Fix nt_extract_public_key for known non-standard contracts

rust/src/core/ton_wallet/mod.rs

@@ -536,7 +536,8 @@ pub unsafe extern "C" fn nt_ton_wallet_prepare_deploy_with_multiple_owners(
                 .handle_error()?
                 .into_iter()
                 .map(parse_public_key)
-                .collect::<Result<Vec<_>, anyhow::Error>>().handle_error()?;
+                .collect::<Result<Vec<_>, anyhow::Error>>()
+                .handle_error()?;
 
             let unsigned_message = ton_wallet
                 .prepare_deploy_with_multiple_owners(expiration, &custodians, req_confirms, None)
@@ -1058,21 +1059,18 @@ pub unsafe extern "C" fn nt_get_wallet_custodians(
 ffi_box!(ton_wallet, Arc<RwLock<TonWallet>>);
 
 #[cfg(test)]
+#[allow(clippy::disallowed_methods)]
 mod test {
-    use std::sync::{Arc, RwLock};
-
-    use nekoton::core::ton_wallet::TonWallet;
+    use std::sync::RwLock;
 
     #[test]
     fn ask_miri() {
         let ptr = Box::into_raw(Box::new(RwLock::new(String::new())));
 
-        let ton_wallet = unsafe { &*(ptr as *mut RwLock<String>) };
+        let ton_wallet = unsafe { &*ptr };
 
-        let len = ton_wallet.read().unwrap().len();
+        _ = ton_wallet.read().unwrap().len();
 
-        unsafe {
-            Box::from_raw(ptr as *mut RwLock<String>);
-        }
+        drop(unsafe { Box::from_raw(ptr) })
     }
 }

rust/src/helpers/abi/mod.rs

@@ -124,7 +124,11 @@ pub unsafe extern "C" fn nt_get_expected_address(
     ) -> Result<serde_json::Value, String> {
         let mut state_init = ton_block::StateInit::construct_from_base64(&tvc).handle_error()?;
         let contract_abi = parse_contract_abi(&contract_abi)?;
-        let public_key = public_key.as_deref().map(parse_public_key).transpose().handle_error()?;
+        let public_key = public_key
+            .as_deref()
+            .map(parse_public_key)
+            .transpose()
+            .handle_error()?;
 
         let params = contract_abi
             .data
@@ -235,7 +239,7 @@ pub unsafe extern "C" fn nt_create_external_message_without_signature(
 
         let time = SystemTime::now()
             .duration_since(UNIX_EPOCH)
-            .map_err(|e| format!("{}", e))?
+            .map_err(|e| e.to_string())?
             .as_millis() as u64;
 
         let expire_at = ExpireAt::new_from_millis(Expiration::Timeout(timeout), time);

rust/src/helpers/mod.rs

@@ -83,12 +83,54 @@ pub unsafe extern "C" fn nt_repack_address(address: *mut c_char) -> *mut c_char
 
 #[no_mangle]
 pub unsafe extern "C" fn nt_extract_public_key(boc: *mut c_char) -> *mut c_char {
+    use nekoton::core::ton_wallet::{highload_wallet_v2, wallet_v3};
+
     let boc = boc.to_string_from_ptr();
 
     fn internal_fn(boc: String) -> Result<serde_json::Value, String> {
-        let public_key = parse_account_stuff(&boc)
-            .and_then(|e| nekoton_abi::extract_public_key(&e).handle_error())
-            .map(hex::encode)?;
+        let public_key = parse_account_stuff(&boc).and_then(|account_stuff| {
+            let state_init = match &account_stuff.storage.state {
+                ton_block::AccountState::AccountActive { state_init, .. } => state_init,
+                _ => return Err(nekoton_abi::ExtractionError::AccountIsNotActive).handle_error(),
+            };
+            let data = match &state_init.data {
+                Some(data) => data,
+                None => {
+                    return Err(nekoton_abi::ExtractionError::AccountDataNotFound).handle_error()
+                },
+            };
+
+            if let Some(code) = &state_init.code {
+                let code_hash = code.repr_hash();
+                if wallet_v3::is_wallet_v3(&code_hash) {
+                    return wallet_v3::InitData::try_from(data).handle_error().and_then(
+                        |init_data| {
+                            ed25519_dalek::PublicKey::from_bytes(init_data.public_key.as_slice())
+                                .map(hex::encode)
+                                .handle_error()
+                        },
+                    );
+                } else if highload_wallet_v2::is_highload_wallet_v2(&code_hash) {
+                    return highload_wallet_v2::InitData::try_from(data)
+                        .handle_error()
+                        .and_then(|init_data| {
+                            ed25519_dalek::PublicKey::from_bytes(init_data.public_key.as_slice())
+                                .map(hex::encode)
+                                .handle_error()
+                        });
+                }
+            }
+
+            let data = ton_types::SliceData::from(data)
+                .get_next_bytes(32)
+                .map_err(|_| nekoton_abi::ExtractionError::CellUnderflow)
+                .handle_error()?;
+
+            ed25519_dalek::PublicKey::from_bytes(&data)
+                .map(hex::encode)
+                .map_err(|_| nekoton_abi::ExtractionError::InvalidPublicKey)
+                .handle_error()
+        })?;
 
         serde_json::to_value(public_key).handle_error()
     }