apollo_starknet_os_program: add security checks for virtual OS

Author: Yoni-Starkware

crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/execution/execute_transactions.cairo

@@ -66,6 +66,7 @@ from starkware.starknet.core.os.execution.execute_transaction_utils import (
     update_class_hash_in_execution_context,
 )
 from starkware.starknet.core.os.execution.revert import init_revert_log
+from starkware.starknet.core.os.execution.execution_constraints import check_n_txs, check_tx_type
 from starkware.starknet.core.os.output import (
     MessageToL2Header,
     OsCarriedOutputs,
@@ -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);

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 ();
+}

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 ();
+}

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"
 }

crates/apollo_starknet_os_program/src/virtual_os_test.rs

@@ -5,7 +5,9 @@ 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"]
+    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"));
 }
 
@@ -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());
 }

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");
+}

crates/starknet_os_flow_tests/src/virtual_os_test_manager.rs

@@ -53,4 +53,16 @@ impl<S: FlowTestState> TestRunner<S> {
     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
+        );
+    }
 }