From 2ed4bc66126513c098da8efe9ab38db1f3f2d129 Mon Sep 17 00:00:00 2001 From: Yonatan Iluz Date: Tue, 6 Jan 2026 20:56:07 +0200 Subject: [PATCH] apollo_starknet_os_program: add security checks for virtual OS --- .../os/execution/execute_transactions.cairo | 6 ++ .../os/execution/execution_constraints.cairo | 10 +++ .../execution_constraints__virtual.cairo | 17 ++++++ .../src/program_hash.json | 4 +- .../src/virtual_os_test.rs | 14 +++-- .../src/virtual_os_test.rs | 61 ++++++++++++++++++- .../src/virtual_os_test_manager.rs | 12 ++++ 7 files changed, 114 insertions(+), 10 deletions(-) create mode 100644 crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints.cairo create mode 100644 crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints__virtual.cairo diff --git a/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execute_transactions.cairo b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execute_transactions.cairo index 0baf5c879f6..86ff33a018b 100644 --- a/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execute_transactions.cairo +++ b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execute_transactions.cairo @@ -65,6 +65,7 @@ from starkware.starknet.core.os.execution.execute_transaction_utils import ( run_validate, update_class_hash_in_execution_context, ) +from starkware.starknet.core.os.execution.execution_constraints import check_n_txs, check_tx_type from starkware.starknet.core.os.execution.revert import init_revert_log from starkware.starknet.core.os.output import ( MessageToL2Header, @@ -144,6 +145,9 @@ func execute_transactions{ local n_txs; %{ OsInputTransactions %} %{ EnterScopeExecuteTransactionsInner %} + + check_n_txs(n_txs=n_txs); + execute_transactions_inner{ builtin_ptrs=builtin_ptrs, contract_state_changes=contract_state_changes, @@ -206,6 +210,8 @@ func execute_transactions_inner{ // Guess the current transaction's type. %{ LoadNextTx %} + check_tx_type(tx_type=tx_type); + if (tx_type == 'INVOKE_FUNCTION') { // Handle the invoke-function transaction. execute_invoke_function_transaction(block_context=block_context); diff --git a/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints.cairo b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints.cairo new file mode 100644 index 00000000000..cd037aee77f --- /dev/null +++ b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints.cairo @@ -0,0 +1,10 @@ +// Execution constraints for transaction execution. +// These are no-op implementations for the Sequencer OS. + +func check_n_txs(n_txs: felt) { + return (); +} + +func check_tx_type(tx_type: felt) { + return (); +} diff --git a/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints__virtual.cairo b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints__virtual.cairo new file mode 100644 index 00000000000..6d2274e525e --- /dev/null +++ b/crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execution_constraints__virtual.cairo @@ -0,0 +1,17 @@ +// Execution constraints for transaction execution (virtual OS version). + +// Checks that the number of transactions is one. +func check_n_txs(n_txs: felt) { + with_attr error_message("Expected exactly one transaction") { + assert n_txs = 1; + } + return (); +} + +// Checks that the transaction type is INVOKE_FUNCTION. +func check_tx_type(tx_type: felt) { + with_attr error_message("Expected INVOKE_FUNCTION transaction") { + assert tx_type = 'INVOKE_FUNCTION'; + } + return (); +} diff --git a/crates/apollo_starknet_os_program/src/program_hash.json b/crates/apollo_starknet_os_program/src/program_hash.json index 813c6adb37e..d02cb566064 100644 --- a/crates/apollo_starknet_os_program/src/program_hash.json +++ b/crates/apollo_starknet_os_program/src/program_hash.json @@ -1,6 +1,6 @@ { - "os": "0x9b874b12cf82d72e25947a557a77a581acb240348275db64e70947e63c1478", - "virtual_os": "0x5c40917ea9e822a504ad48f7a57f2485a07d5cdfe1e8b1b38b7ca6627d8daf4", + "os": "0x78179950037794503935395abfcc7446fbdee452d388a9cc047ef7201bc2278", + "virtual_os": "0x6c33b6ac546a56d5f79fe7367a3c4ec550ba8bf627d97f5cfe0f6a6cf120e5c", "aggregator": "0x63868c8dbf389119d95bae88b80043f5d237d3ddec072fb8f2095fcfacb9b1c", "aggregator_with_prefix": "0xd892fd2fd978d8c2749a6678457ca99161f3d1822e4c3b8c03914817c6de1a" } \ No newline at end of file diff --git a/crates/apollo_starknet_os_program/src/virtual_os_test.rs b/crates/apollo_starknet_os_program/src/virtual_os_test.rs index ff1c75b41b8..4155ae7656b 100644 --- a/crates/apollo_starknet_os_program/src/virtual_os_test.rs +++ b/crates/apollo_starknet_os_program/src/virtual_os_test.rs @@ -5,8 +5,10 @@ use crate::{OS_PROGRAM, VIRTUAL_OS_PROGRAM, VIRTUAL_OS_SWAPPED_FILES}; /// Asserts the list of swapped virtual OS files matches the expected list. #[test] fn test_virtual_os_swapped_files() { - expect!["starkware/starknet/core/os/os_utils.cairo"] - .assert_eq(&VIRTUAL_OS_SWAPPED_FILES.join("\n")); + expect![[r#" + starkware/starknet/core/os/execution/execution_constraints.cairo + starkware/starknet/core/os/os_utils.cairo"#]] + .assert_eq(&VIRTUAL_OS_SWAPPED_FILES.join("\n")); } /// Asserts the bytecode length of the OS program and virtual OS program match expected values. @@ -14,11 +16,11 @@ fn test_virtual_os_swapped_files() { #[test] fn test_program_bytecode_lengths() { expect![[r#" - 15527 -"#]] + 15535 + "#]] .assert_debug_eq(&OS_PROGRAM.data_len()); expect![[r#" - 13172 -"#]] + 13184 + "#]] .assert_debug_eq(&VIRTUAL_OS_PROGRAM.data_len()); } diff --git a/crates/starknet_os_flow_tests/src/virtual_os_test.rs b/crates/starknet_os_flow_tests/src/virtual_os_test.rs index c2af4066f35..928aa60c3c5 100644 --- a/crates/starknet_os_flow_tests/src/virtual_os_test.rs +++ b/crates/starknet_os_flow_tests/src/virtual_os_test.rs @@ -2,8 +2,11 @@ use blockifier_test_utils::cairo_versions::{CairoVersion, RunnableCairo1}; use blockifier_test_utils::calldata::create_calldata; use blockifier_test_utils::contracts::FeatureContract; use rstest::rstest; -use starknet_api::core::EthAddress; -use starknet_api::transaction::{L2ToL1Payload, MessageToL1}; +use starknet_api::abi::abi_utils::selector_from_name; +use starknet_api::core::{EthAddress, Nonce}; +use starknet_api::executable_transaction::L1HandlerTransaction as ExecutableL1HandlerTransaction; +use starknet_api::transaction::fields::Fee; +use starknet_api::transaction::{L1HandlerTransaction, L2ToL1Payload, MessageToL1}; use starknet_api::{calldata, invoke_tx_args}; use starknet_types_core::felt::Felt; @@ -34,3 +37,57 @@ async fn test_basic_happy_flow() { test_builder.build().await.run_virtual_and_validate(); } + +/// Security tests. +/// Note that it's important to construct the hints correctly and get the error directly from Cairo +/// (and not from the blockifier), as users can submit virtual OS proofs with arbitrary hints. + +#[rstest] +#[tokio::test] +/// Test that the virtual OS fails when more than one transaction is added. +async fn test_two_txs_os_error() { + let test_contract = FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm)); + + let (mut test_builder, [contract_address]) = + TestBuilder::create_standard_virtual([(test_contract, calldata![Felt::ONE, Felt::TWO])]) + .await; + + // Add first invoke transaction. + let calldata = create_calldata(contract_address, "test_storage_read", &[Felt::ONE]); + test_builder.add_funded_account_invoke(invoke_tx_args! { calldata: calldata.clone() }); + + // Add second invoke transaction - this should cause the virtual OS to fail. + test_builder.add_funded_account_invoke(invoke_tx_args! { calldata }); + + test_builder.build().await.run_virtual_expect_error("Expected exactly one transaction"); +} + +#[rstest] +#[tokio::test] +/// Test that the virtual OS fails when a non-invoke transaction is added. +async fn test_non_invoke_tx_os_error() { + let test_contract = FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm)); + + let (mut test_builder, [contract_address]) = + TestBuilder::create_standard_virtual([(test_contract, calldata![Felt::ONE, Felt::TWO])]) + .await; + + // Add an L1 handler transaction instead of an invoke. + // TODO(Yoni): parameterize other transaction types. + let tx = ExecutableL1HandlerTransaction::create( + L1HandlerTransaction { + version: L1HandlerTransaction::VERSION, + nonce: Nonce::default(), + contract_address, + entry_point_selector: selector_from_name("l1_handle"), + // from_address, arg. + calldata: calldata![Felt::ONE, Felt::TWO], + }, + &test_builder.chain_id(), + Fee(1_000_000), + ) + .unwrap(); + test_builder.add_l1_handler_tx(tx, None); + + test_builder.build().await.run_virtual_expect_error("Expected INVOKE_FUNCTION transaction"); +} diff --git a/crates/starknet_os_flow_tests/src/virtual_os_test_manager.rs b/crates/starknet_os_flow_tests/src/virtual_os_test_manager.rs index 8de6ba25809..789cec27a19 100644 --- a/crates/starknet_os_flow_tests/src/virtual_os_test_manager.rs +++ b/crates/starknet_os_flow_tests/src/virtual_os_test_manager.rs @@ -53,4 +53,16 @@ impl TestRunner { pub(crate) fn run_virtual_and_validate(self) { self.run_virtual().validate(); } + + /// Runs the virtual OS and expects it to fail with an error containing the given string. + pub(crate) fn run_virtual_expect_error(self, expected_error: &str) { + let err = run_virtual_os(self.os_hints).expect_err("Expected virtual OS to fail"); + let err_string = err.to_string(); + assert!( + err_string.contains(expected_error), + "Expected error to contain '{}', got: {}", + expected_error, + err_string + ); + } }