From 5c202227272435fccb184a8fa9dd05be2a988323 Mon Sep 17 00:00:00 2001
From: OuyangHang33 <hank.ouyang@intel.com>
Date: Sun, 28 Apr 2024 16:54:06 +0800
Subject: [PATCH] Add rollback test for spdm-rs and spdm-emu

Signed-off-by: OuyangHang33 <hank.ouyang@intel.com>
---
 .github/workflows/main.yml          |   8 ++
 sh_script/rollbacktest.sh           | 127 +++++++++++++++++++
 test/spdm-requester-emu/Cargo.toml  |   2 +
 test/spdm-requester-emu/src/main.rs | 181 ++++++++++++++++++++++++++++
 4 files changed, 318 insertions(+)
 create mode 100755 sh_script/rollbacktest.sh

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 504569a..6570aad 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -89,3 +89,11 @@ jobs:
           RUN_RESPONDER_FEATURES: ${{ matrix.run_responder_features }}
         run: |
           ./sh_script/build.sh -r
+
+      - name: Rollback test
+        shell: bash
+        env:
+          CC_x86_64_unknown_none: clang
+          AR_x86_64_unknown_none: llvm-ar
+        run: |
+          ./sh_script/rollbacktest.sh -a
diff --git a/sh_script/rollbacktest.sh b/sh_script/rollbacktest.sh
new file mode 100755
index 0000000..fad5ab6
--- /dev/null
+++ b/sh_script/rollbacktest.sh
@@ -0,0 +1,127 @@
+#!/bin/bash
+
+set -euo pipefail
+
+export RUST_MIN_STACK=10485760
+RUSTFLAGS=${RUSTFLAGS:-}
+
+usage() {
+    cat <<EOM
+Usage: $(basename "$0") [OPTION]...
+  -e EMU test
+  -r Rust test
+  -a Run all roll back test
+  -h Show help info
+EOM
+}
+
+trap cleanup exit
+
+echo_command() {
+    set -x
+    "$@"
+    set +x
+}
+
+cleanup() {
+    kill -9 $(ps aux | grep spdm-responder | grep emu | awk '{print $2}') || true
+    kill -9 $(ps aux | grep spdm_responder_emu | grep emu | awk '{print $2}') || true
+}
+
+RUN_REQUESTER_FEATURES=${RUN_REQUESTER_FEATURES:-spdm-ring,hashed-transcript-data,async-executor}
+RUN_RESPONDER_FEATURES=${RUN_RESPONDER_FEATURES:-spdm-ring,hashed-transcript-data,async-executor}
+RUN_REQUESTER_UPDATE_KEYS_FEATURES="${RUN_REQUESTER_FEATURES},test_update_keys"
+RUN_REQUESTER_VERIFY_KEYS_FEATURES="${RUN_REQUESTER_FEATURES},test_verify_keys"
+
+
+run_with_spdm_emu_rsp_rust_req() {
+    echo "Running spdm-rs requester to test spdm-emu responder..."
+    pushd test_key
+    chmod +x ./spdm_responder_emu
+    echo_command  ./spdm_responder_emu --trans PCI_DOE &
+    popd
+    sleep 5
+    echo_command cargo run -p spdm-requester-emu --no-default-features --features="$RUN_REQUESTER_UPDATE_KEYS_FEATURES"
+    cleanup    
+    popd
+}
+
+run_with_spdm_emu_rsp_rust_req_verify() {
+    echo "Running spdm-rs requester to test spdm-emu responder..."
+    pushd test_key
+    chmod +x ./spdm_responder_emu
+    echo_command  ./spdm_responder_emu --trans PCI_DOE &
+    popd
+    sleep 5
+    echo_command cargo run -p spdm-requester-emu --no-default-features --features="$RUN_REQUESTER_VERIFY_KEYS_FEATURES"
+    cleanup    
+    popd
+}
+
+run_rust_spdm_emu_test_rsp() {
+    echo "Running spdm-rs requester to test spdm-rs responder..."
+    echo_command cargo run -p spdm-responder-emu --no-default-features --features="$RUN_RESPONDER_FEATURES" &
+    sleep 5
+    echo_command cargo run -p spdm-requester-emu --no-default-features --features="$RUN_REQUESTER_UPDATE_KEYS_FEATURES"
+    cleanup
+}
+
+run_rust_spdm_emu_test_rsp_verify() {
+    echo "Running spdm-rs requester to test spdm-rs responder..."
+    echo_command cargo run -p spdm-responder-emu --no-default-features --features="$RUN_RESPONDER_FEATURES" &
+    sleep 5
+    echo_command cargo run -p spdm-requester-emu --no-default-features --features="$RUN_REQUESTER_VERIFY_KEYS_FEATURES"
+    cleanup
+}
+
+EMU_OPTION=false
+RUST_OPTION=false
+ALL_OPTION=false
+
+
+process_args() {
+    while getopts ":erah" option; do
+        case "${option}" in
+            e)
+                EMU_OPTION=true
+            ;;
+            r)
+                RUST_OPTION=true
+            ;;
+            a)
+                ALL_OPTION=true
+            ;;
+            h)
+                usage
+                exit 0
+            ;;
+            *)
+                echo "Invalid option '-$OPTARG'"
+                usage
+                exit 1
+            ;;
+        esac
+    done
+}
+
+
+main() {
+    if [[ ${EMU_OPTION} == true ]]; then
+        run_with_spdm_emu_rsp_rust_req
+        run_with_spdm_emu_rsp_rust_req_verify
+    fi
+    if [[ ${RUST_OPTION} == true ]]; then
+        run_rust_spdm_emu_test_rsp
+        run_rust_spdm_emu_test_rsp_verify
+    fi
+    if [[ ${ALL_OPTION} == true ]]; then
+        run_with_spdm_emu_rsp_rust_req
+        run_with_spdm_emu_rsp_rust_req_verify
+        run_rust_spdm_emu_test_rsp
+        run_rust_spdm_emu_test_rsp_verify
+
+    fi
+}
+
+process_args "$@"
+main
\ No newline at end of file
diff --git a/test/spdm-requester-emu/Cargo.toml b/test/spdm-requester-emu/Cargo.toml
index 84ef2fb..d93ed18 100644
--- a/test/spdm-requester-emu/Cargo.toml
+++ b/test/spdm-requester-emu/Cargo.toml
@@ -35,3 +35,5 @@ async-tokio = ["spdm-emu/async-tokio"]
 is_sync = ["spdm-emu/is_sync", "spdmlib/is_sync", "maybe-async/is_sync", "idekm/is_sync", "tdisp/is_sync", "mctp_transport/is_sync", "pcidoe_transport/is_sync"]
 test_stack_size = ["td-benchmark"]
 test_heap_size = ["dhat"]
+test_verify_keys = []
+test_update_keys = []
diff --git a/test/spdm-requester-emu/src/main.rs b/test/spdm-requester-emu/src/main.rs
index 219ae6f..9793e70 100644
--- a/test/spdm-requester-emu/src/main.rs
+++ b/test/spdm-requester-emu/src/main.rs
@@ -354,6 +354,187 @@ async fn test_spdm(
             panic!("send_receive_spdm_heartbeat failed");
         }
 
+        #[cfg(feature = "test_update_keys")]
+        {
+            context.common.reset_buffer_via_request_code(
+                SpdmRequestResponseCode::SpdmRequestKeyUpdate,
+                Some(session_id),
+            );
+
+            // Get original keys.
+            let ori_keys = &context.common.session[0].export_keys();
+
+            // Update all keys and send wrong UpdateAllKeys request.
+            let mut send_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context.encode_spdm_key_update_op(
+                SpdmKeyUpdateOperation::Unknown(0),
+                1,
+                &mut send_buffer,
+            ) {
+                u
+            } else {
+                panic!("Send request fail!\n");
+            };
+            let _ = context
+                .send_message(Some(session_id), &send_buffer[..used], false)
+                .await
+                .is_ok();
+
+            let session = if let Some(s) = context.common.get_session_via_id(session_id) {
+                s
+            } else {
+                panic!("Invalid session id!\n");
+            };
+
+            // Create new keys after sending KeyUpdate request.
+            let _ = session
+                .create_data_secret_update(SpdmVersion::SpdmVersion12, true, true)
+                .ok();
+
+            // Update all keys to new keys.
+            let new_keys = &context.common.session[0].export_keys();
+            assert!(new_keys.0.encryption_key.data != ori_keys.0.encryption_key.data);
+            assert!(new_keys.1.encryption_key.data != ori_keys.1.encryption_key.data);
+
+            let mut receive_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context
+                .receive_message(Some(session_id), &mut receive_buffer, false)
+                .await
+            {
+                u
+            } else {
+                panic!("Receive message fail!\n");
+            };
+
+            let status = context.handle_spdm_key_update_op_response(
+                session_id,
+                true,
+                true,
+                &receive_buffer[..used],
+            );
+            // Return key update fail due to wrong request message
+            assert!(status.is_err());
+            let rollbacked_keys = &context.common.session[0].export_keys();
+            // Rollback all keys to original keys.
+            assert!(rollbacked_keys.0.encryption_key.data == ori_keys.0.encryption_key.data);
+            assert!(rollbacked_keys.1.encryption_key.data == ori_keys.1.encryption_key.data);
+        }
+
+        #[cfg(feature = "test_verify_keys")]
+        {
+            context.common.reset_buffer_via_request_code(
+                SpdmRequestResponseCode::SpdmRequestKeyUpdate,
+                Some(session_id),
+            );
+
+            // Get original keys.
+            let ori_keys = &context.common.session[0].export_keys();
+
+            // 1. Send UpdateAllKeys request and update all keys.
+            let mut send_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context.encode_spdm_key_update_op(
+                SpdmKeyUpdateOperation::SpdmUpdateAllKeys,
+                1,
+                &mut send_buffer,
+            ) {
+                u
+            } else {
+                panic!("Send request fail!\n");
+            };
+            let _ = context
+                .send_message(Some(session_id), &send_buffer[..used], false)
+                .await
+                .is_ok();
+
+            let session = if let Some(s) = context.common.get_session_via_id(session_id) {
+                s
+            } else {
+                panic!("Invalid session id!\n");
+            };
+
+            // Create new keys after sending KeyUpdate request.
+            let _ = session
+                .create_data_secret_update(SpdmVersion::SpdmVersion12, true, true)
+                .ok();
+
+            // Update all keys to new keys.
+            let new_keys = &context.common.session[0].export_keys();
+            assert!(new_keys.0.encryption_key.data != ori_keys.0.encryption_key.data);
+
+            let mut receive_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context
+                .receive_message(Some(session_id), &mut receive_buffer, false)
+                .await
+            {
+                u
+            } else {
+                panic!("Receive message fail!\n");
+            };
+
+            let status = context.handle_spdm_key_update_op_response(
+                session_id,
+                true,
+                true,
+                &receive_buffer[..used],
+            );
+            // Key update successfully.
+            assert!(status.is_ok());
+            let new_keys = &context.common.session[0].export_keys();
+            assert!(new_keys.0.encryption_key.data != ori_keys.0.encryption_key.data);
+            assert!(new_keys.1.encryption_key.data != ori_keys.1.encryption_key.data);
+
+            // 2. Send VerifyNewKeys request.
+            context.common.reset_buffer_via_request_code(
+                SpdmRequestResponseCode::SpdmRequestKeyUpdate,
+                Some(session_id),
+            );
+            let mut send_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context.encode_spdm_key_update_op(
+                SpdmKeyUpdateOperation::Unknown(0),
+                2,
+                &mut send_buffer,
+            ) {
+                u
+            } else {
+                panic!("Send request fail!\n");
+            };
+            let _ = context
+                .send_message(Some(session_id), &send_buffer[..used], false)
+                .await
+                .is_ok();
+
+            let session = if let Some(s) = context.common.get_session_via_id(session_id) {
+                s
+            } else {
+                panic!("Invalid session id!\n");
+            };
+
+            // No keys update in VerifyNewKeys step.
+            let _ = session
+                .create_data_secret_update(SpdmVersion::SpdmVersion12, false, false)
+                .ok();
+
+            // Receive verify response message.
+            let mut receive_buffer = [0u8; config::MAX_SPDM_MSG_SIZE];
+            let used = if let Ok(u) = context
+                .receive_message(Some(session_id), &mut receive_buffer, false)
+                .await
+            {
+                u
+            } else {
+                panic!("Receive message fail!\n");
+            };
+
+            let status = context.handle_spdm_key_update_op_response(
+                session_id,
+                false,
+                false,
+                &receive_buffer[..used],
+            );
+            // Return VerifyNewKeys fail due to wrong request message
+            assert!(status.is_err());
+        }
+
         if context
             .send_receive_spdm_key_update(session_id, SpdmKeyUpdateOperation::SpdmUpdateAllKeys)
             .await