Skip to content

Commit 9f39d8a

Browse files
Kailai-WangKailai-Wang
authored
Fix missing types in manual tests (#3634)
* update types * fix tests * cargo update * update names
1 parent 5fec0ed commit 9f39d8a

File tree

9 files changed

+309
-245
lines changed

9 files changed

+309
-245
lines changed

‎tee-worker/omni-executor/Cargo.lock‎

Lines changed: 227 additions & 223 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tee-worker/omni-executor/aa-contracts-client/src/entry_point_client.rs‎

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
use crate::types::{
1818
createAccountCall, depositToCall, getSenderAddressCall, getUserOpHashCall, handleOpsCall,
19-
simulateHandleOpsCall, simulateValidationCall, ExecutionResult, SenderAddressResult,
19+
simulateHandleOpsCall, simulateValidationCall, ExecutionResult, OwnerType, SenderAddressResult,
2020
ValidationResult,
2121
};
2222
use crate::utils::{
@@ -511,10 +511,18 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
511511
factory_address: Address,
512512
account_implementation: Address,
513513
oa: FixedBytes<32>,
514+
oa_type: OwnerType,
514515
client_id: &[u8],
515516
root: Address,
516517
) -> Address {
517-
calculate_omni_account_address(factory_address, account_implementation, oa, client_id, root)
518+
calculate_omni_account_address(
519+
factory_address,
520+
account_implementation,
521+
oa,
522+
oa_type,
523+
client_id,
524+
root,
525+
)
518526
}
519527

520528
pub async fn deposit_to(&self, account: Address, amount: U256) -> Result<String, ()> {
@@ -542,14 +550,15 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
542550
&self,
543551
factory_address: Address,
544552
oa: [u8; 32],
553+
oa_type: OwnerType,
545554
client_id: &[u8],
546555
root_address: Address,
547556
call_data: Bytes,
548557
paymaster_address: Option<Address>,
549558
) -> Result<PackedUserOperation, ()> {
550559
// Create init code using existing helper
551560
let init_code_bytes =
552-
prepare_factory_init_code(factory_address, oa, client_id, root_address);
561+
prepare_factory_init_code(factory_address, oa, oa_type, client_id, root_address);
553562
let init_code = Bytes::from(init_code_bytes);
554563

555564
// Get sender address from EntryPoint
@@ -560,6 +569,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
560569
sender,
561570
factory_address,
562571
oa,
572+
oa_type,
563573
client_id,
564574
root_address,
565575
call_data,
@@ -576,6 +586,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
576586
factory_address: Address,
577587
account_implementation: Address,
578588
oa: [u8; 32],
589+
oa_type: OwnerType,
579590
client_id: &[u8],
580591
root_address: Address,
581592
call_data: Bytes,
@@ -588,6 +599,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
588599
factory_address,
589600
account_implementation,
590601
oa_fixed,
602+
oa_type,
591603
client_id,
592604
root_address,
593605
);
@@ -597,6 +609,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
597609
sender,
598610
factory_address,
599611
oa,
612+
oa_type,
600613
client_id,
601614
root_address,
602615
call_data,
@@ -612,6 +625,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
612625
sender: Address,
613626
factory_address: Address,
614627
oa: [u8; 32],
628+
oa_type: OwnerType,
615629
client_id: &[u8],
616630
root_address: Address,
617631
call_data: Bytes,
@@ -633,7 +647,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
633647
let init_code_to_use = if code.is_empty() {
634648
// No code at address, include init code
635649
let init_code_bytes =
636-
prepare_factory_init_code(factory_address, oa, client_id, root_address);
650+
prepare_factory_init_code(factory_address, oa, oa_type, client_id, root_address);
637651
Bytes::from(init_code_bytes)
638652
} else {
639653
// Code already exists, no init code needed
@@ -675,6 +689,7 @@ impl<P: RpcProvider<Transaction = TransactionRequest, Addr = Address>> EntryPoin
675689
pub fn prepare_factory_init_code(
676690
factory_address: Address,
677691
oa: [u8; 32],
692+
oa_type: OwnerType,
678693
client_id: &[u8],
679694
root: Address,
680695
) -> Vec<u8> {
@@ -683,6 +698,7 @@ pub fn prepare_factory_init_code(
683698
let mut call = createAccountCall {
684699
//safe to unwrap
685700
oa: oa.into(),
701+
oaType: oa_type,
686702
//safe to unwrap
687703
clientId: client_id.to_owned().into(),
688704
root,
@@ -729,7 +745,7 @@ pub fn create_gas_fees(max_fee_per_gas: U256, max_priority_fee_per_gas: U256) ->
729745

730746
#[cfg(test)]
731747
pub mod test {
732-
use crate::types::depositCall;
748+
use crate::types::{depositCall, OwnerType};
733749
use crate::utils::build_payable_transaction;
734750
use crate::{prepare_factory_init_code, EntryPointClient, GasPriceConfig};
735751
use alloy::hex;
@@ -772,6 +788,7 @@ pub mod test {
772788
let init_code_bytes = prepare_factory_init_code(
773789
factory_address,
774790
oa_bytes.0,
791+
OwnerType::Evm,
775792
client_id_fixed_bytes.as_ref(),
776793
root_address,
777794
);
@@ -784,7 +801,7 @@ pub mod test {
784801

785802
#[test(test)]
786803
pub fn test_prepare_init_code() {
787-
let expected = "e7f1725e7734ce288f8367e1bb143e90bb3f05120db21afe6d659c2361df3754aa7d46d9f0c993ec0335e60508128c6e3843f7dd962113910000000000000000000000000000000000000000000000000000000000000060000000000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000000";
804+
let expected = "e7f1725e7734ce288f8367e1bb143e90bb3f0512158b8ca56d659c2361df3754aa7d46d9f0c993ec0335e60508128c6e3843f7dd9621139100000000000000000000000000000000000000000000000000000000000000060000000000000000000000000000000000000000000000000000000000000080000000000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000000";
788805

789806
let factory_address = address!("0xe7f1725E7734CE288F8367e1Bb143E90bb3F0512");
790807
let oa: [u8; 32] =
@@ -797,7 +814,13 @@ pub mod test {
797814
.unwrap();
798815
let root_address = address!("0x0000000000000000000000000000000000000001");
799816

800-
let init_code = prepare_factory_init_code(factory_address, oa, &client_id, root_address);
817+
let init_code = prepare_factory_init_code(
818+
factory_address,
819+
oa,
820+
OwnerType::Evm,
821+
&client_id,
822+
root_address,
823+
);
801824

802825
assert_eq!(expected, hex::encode(init_code));
803826
}
@@ -822,6 +845,7 @@ pub mod test {
822845
let init_code_bytes = prepare_factory_init_code(
823846
factory_address,
824847
oa_bytes.0,
848+
OwnerType::Evm,
825849
&client_id_fixed_bytes,
826850
root_address,
827851
);
@@ -861,6 +885,7 @@ pub mod test {
861885
.create_packed_user_operation(
862886
factory_address,
863887
oa_bytes.0,
888+
OwnerType::Evm,
864889
client_id_bytes,
865890
root_address,
866891
call_data1,
@@ -875,6 +900,7 @@ pub mod test {
875900
.create_packed_user_operation(
876901
factory_address,
877902
oa_bytes.0,
903+
OwnerType::Evm,
878904
client_id_bytes,
879905
root_address,
880906
call_data2,
@@ -914,6 +940,7 @@ pub mod test {
914940
#[test(tokio::test)]
915941
#[ignore = "manual"]
916942
pub async fn try_full_flow() {
943+
let oa_type = OwnerType::Evm;
917944
let user_signer = PrivateKeySigner::from_str(
918945
"0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6",
919946
)
@@ -945,6 +972,7 @@ pub mod test {
945972
let init_code_bytes = prepare_factory_init_code(
946973
factory_address,
947974
oa_bytes.0,
975+
oa_type,
948976
&client_id_fixed_bytes,
949977
root_address,
950978
);
@@ -956,6 +984,7 @@ pub mod test {
956984
.create_packed_user_operation(
957985
factory_address,
958986
oa_bytes.0,
987+
oa_type,
959988
&client_id_fixed_bytes,
960989
root_address,
961990
call_data.clone(),
@@ -1010,6 +1039,7 @@ pub mod test {
10101039
#[ignore = "manual"]
10111040
pub async fn test_local_vs_entrypoint_address_calculation() {
10121041
let client_id = "test_client";
1042+
let oa_type = OwnerType::Evm;
10131043
let user_address = address!("0xa0Ee7A142d267C1f36714E4a8F75612F20a79720");
10141044
let entrypoint_address = address!("0x5FbDB2315678afecb367f032d93F642f64180aa3");
10151045
let factory_address = address!("0xe7f1725E7734CE288F8367e1Bb143E90bb3F0512");
@@ -1026,8 +1056,13 @@ pub mod test {
10261056
let oa_bytes: FixedBytes<32> = FixedBytes::from_slice(oa.as_ref());
10271057

10281058
// Calculate address using EntryPoint.getSenderAddress
1029-
let init_code_bytes =
1030-
prepare_factory_init_code(factory_address, oa_bytes.0, client_id_bytes, root_address);
1059+
let init_code_bytes = prepare_factory_init_code(
1060+
factory_address,
1061+
oa_bytes.0,
1062+
oa_type,
1063+
client_id_bytes,
1064+
root_address,
1065+
);
10311066
let init_code = Bytes::from(init_code_bytes);
10321067
let entrypoint_address_result = entrypoint_client
10331068
.get_sender_address(init_code)
@@ -1039,6 +1074,7 @@ pub mod test {
10391074
factory_address,
10401075
account_implementation,
10411076
oa_bytes,
1077+
oa_type,
10421078
client_id_bytes,
10431079
root_address,
10441080
);

‎tee-worker/omni-executor/aa-contracts-client/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,5 @@ pub use entry_point_client::{
2323
prepare_factory_init_code, EntryPointClient, GasPriceConfig, RetryConfig,
2424
};
2525
pub use omni_account_client::OmniAccountClient;
26-
pub use types::PackedUserOperation;
26+
pub use types::{OwnerType, PackedUserOperation};
2727
pub use utils::{calculate_omni_account_address, calculate_user_operation_hash};

‎tee-worker/omni-executor/aa-contracts-client/src/types.rs‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,21 @@ sol! {
5151
bytes32 paymasterAndData;
5252
}
5353

54+
// Owner type
55+
enum OwnerType {
56+
Pumpx,
57+
Email,
58+
Twitter,
59+
Discord,
60+
Github,
61+
Substrate,
62+
Evm,
63+
Bitcoin,
64+
Solana,
65+
Google,
66+
Passkey
67+
}
68+
5469
// Passkey public key
5570
struct PasskeyPublicKey {
5671
bytes32 x;
@@ -65,14 +80,14 @@ sol! {
6580
error SenderAddressResult(address sender);
6681

6782
// smart account factory
68-
function createAccount(bytes32 oa, bytes memory clientId, address root) public;
69-
function getAddress(bytes32 oa, bytes memory clientId, address root) public view returns (address);
83+
function createAccount(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) public;
84+
function getAddress(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) public view returns (address);
7085

7186
// smart account
7287
function getNonce() public view virtual returns (uint256);
7388
function addRootSigner(address root) public;
7489
function removeRootSigner(address root) public;
75-
function initialize(bytes32 oa, bytes memory clientId, address root) public;
90+
function initialize(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) public;
7691
function getOwner() public view returns (bytes32);
7792

7893
// passkey signer management

‎tee-worker/omni-executor/aa-contracts-client/src/utils.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
// You should have received a copy of the GNU General Public License
1515
// along with Litentry. If not, see <https://www.gnu.org/licenses/>.
1616

17-
use crate::types::PackedUserOperationForHashing;
17+
use crate::types::{OwnerType, PackedUserOperationForHashing};
1818
use crate::PackedUserOperation;
1919
use alloy::primitives::{keccak256, Address, Bytes, FixedBytes, TxKind, U256};
2020
use alloy::rpc::types::{TransactionInput, TransactionRequest};
@@ -49,12 +49,14 @@ pub fn calculate_omni_account_address(
4949
factory_address: Address,
5050
account_implementation: Address,
5151
oa: FixedBytes<32>,
52+
oa_type: OwnerType,
5253
client_id: &[u8],
5354
root: Address,
5455
) -> Address {
5556
use crate::types::initializeCall;
5657

57-
let initialize_call = initializeCall { oa, clientId: Bytes::from(client_id.to_vec()), root };
58+
let initialize_call =
59+
initializeCall { oa, oaType: oa_type, clientId: Bytes::from(client_id.to_vec()), root };
5860
let initialize_data = initialize_call.abi_encode();
5961

6062
let mut constructor_params = Vec::new();
@@ -278,6 +280,7 @@ mod tests {
278280
factory_address,
279281
account_implementation,
280282
oa,
283+
OwnerType::Evm,
281284
client_id,
282285
root,
283286
);
@@ -288,6 +291,7 @@ mod tests {
288291
factory_address,
289292
account_implementation,
290293
oa,
294+
OwnerType::Evm,
291295
client_id,
292296
root,
293297
);
@@ -312,6 +316,7 @@ mod tests {
312316
factory_address,
313317
account_implementation,
314318
oa,
319+
OwnerType::Evm,
315320
client_id_1,
316321
root,
317322
);
@@ -320,6 +325,7 @@ mod tests {
320325
factory_address,
321326
account_implementation,
322327
oa,
328+
OwnerType::Evm,
323329
client_id_2,
324330
root,
325331
);

‎tee-worker/omni-executor/aa-contracts/deploy-docker.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ deploy_contracts() {
6969
BROADCAST_FILE="$SCRIPT_DIR/broadcast/DeployLocal.s.sol/$CHAIN_ID/run-latest.json"
7070

7171
if [ -f "$BROADCAST_FILE" ]; then
72-
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPoint"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
73-
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactory"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
72+
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPointV1"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
73+
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactoryV1"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
7474
PAYMASTER_ADDRESS=$(grep -A2 '"contractName": "SimplePaymaster"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
7575

7676
# Extract test token addresses

‎tee-worker/omni-executor/aa-contracts/deploy-local.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,8 @@ deploy_contracts() {
8080
BROADCAST_FILE="$SCRIPT_DIR/broadcast/DeployLocalWithPaymaster.s.sol/$CHAIN_ID/run-latest.json"
8181

8282
if [ -f "$BROADCAST_FILE" ]; then
83-
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPoint"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
84-
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactory"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
83+
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPointV1"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
84+
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactoryV1"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
8585
# Try to find DemoPaymaster first, fall back to SimplePaymaster
8686
PAYMASTER_ADDRESS=$(grep -A2 '"contractName": "DemoPaymaster"' "$BROADCAST_FILE" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
8787
if [ -z "$PAYMASTER_ADDRESS" ]; then

‎tee-worker/omni-executor/aa-contracts/update-demo-addresses.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,9 @@ extract_addresses() {
7575
fi
7676

7777
# Extract addresses using grep and sed
78-
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPoint"' "$broadcast_file" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
78+
ENTRYPOINT_ADDRESS=$(grep -A2 '"contractName": "EntryPointV1"' "$broadcast_file" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
7979
# Try both SmartAccountFactory (old) and OmniAccountFactory (new) names
80-
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactory"' "$broadcast_file" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
80+
FACTORY_ADDRESS=$(grep -A2 '"contractName": "OmniAccountFactoryV1"' "$broadcast_file" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
8181
if [ -z "$FACTORY_ADDRESS" ]; then
8282
FACTORY_ADDRESS=$(grep -A2 '"contractName": "SmartAccountFactory"' "$broadcast_file" | grep '"contractAddress"' | sed 's/.*"contractAddress": "\(.*\)".*/\1/' | head -1)
8383
fi

0 commit comments

Comments
 (0)