// Copyright (c) 2020 Intel Corporation

//

// SPDX-License-Identifier: Apache-2.0 or MIT

use fuzzlib::{

    spdmlib::{

        common::session::{SpdmSession, SpdmSessionState},

        message::SpdmKeyExchangeMutAuthAttributes,

    },

    *,

};

use spdmlib::protocol::*;

use spin::Mutex;

extern crate alloc;

use alloc::sync::Arc;

async fn fuzz_send_receive_spdm_finish(fuzzdata: Arc<Vec<u8>>) {

    spdmlib::secret::asym_sign::register(SECRET_ASYM_IMPL_INSTANCE.clone());

    spdmlib::secret::pqc_asym_sign::register(SECRET_PQC_ASYM_IMPL_INSTANCE.clone());

    spdmlib::crypto::aead::register(FAKE_AEAD.clone());

    spdmlib::crypto::hmac::register(FAKE_HMAC.clone());

    spdmlib::crypto::hkdf::register(FAKE_HKDF.clone());

    // TCD:

    // - id: 0

    // - title: 'Fuzz SPDM handle finish response'

    // - description: '<p>Request finish to complete the handshake, and the handshake is performed in the clear.</p>'

    // -

    {

        let (req_config_info, req_provision_info) = req_create_info();

        let shared_buffer = SharedBuffer::new();

        let pcidoe_transport_encap = Arc::new(Mutex::new(PciDoeTransportEncap {}));

        let mut device_io_requester =

            fake_device_io::FakeSpdmDeviceIo::new(Arc::new(shared_buffer));

        device_io_requester.set_rx(&fuzzdata);

        let device_io_requester = Arc::new(Mutex::new(device_io_requester));

        let mut requester = requester::RequesterContext::new(

            device_io_requester,

            pcidoe_transport_encap,

            req_config_info,

            req_provision_info,

        );

        requester.common.negotiate_info.spdm_version_sel = SpdmVersion::SpdmVersion12;

        requester.common.negotiate_info.req_ct_exponent_sel = 0;

        requester.common.negotiate_info.req_capabilities_sel = SpdmRequestCapabilityFlags::CERT_CAP

            | SpdmRequestCapabilityFlags::HANDSHAKE_IN_THE_CLEAR_CAP;

        requester.common.negotiate_info.rsp_ct_exponent_sel = 0;

        requester.common.negotiate_info.rsp_capabilities_sel = SpdmResponseCapabilityFlags::CERT_CAP

            | SpdmResponseCapabilityFlags::HANDSHAKE_IN_THE_CLEAR_CAP;

        requester.common.negotiate_info.base_hash_sel = SpdmBaseHashAlgo::TPM_ALG_SHA_384;

        requester.common.negotiate_info.base_asym_sel =

            SpdmBaseAsymAlgo::TPM_ALG_ECDSA_ECC_NIST_P384;

        requester.common.negotiate_info.dhe_sel = SpdmDheAlgo::SECP_384_R1;

        requester.common.negotiate_info.aead_sel = SpdmAeadAlgo::AES_256_GCM;

        requester.common.negotiate_info.req_asym_sel = SpdmReqAsymAlgo::TPM_ALG_RSAPSS_2048;

        requester.common.negotiate_info.key_schedule_sel = SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE;

        requester.common.reset_runtime_info();

        requester.common.session[0] = SpdmSession::new();

        requester.common.session[0].setup(4294836221).unwrap();

        requester.common.session[0].set_crypto_param(

            SpdmBaseHashAlgo::TPM_ALG_SHA_384,

            SpdmDheAlgo::SECP_384_R1,

            SpdmKemAlgo::empty(),

            SpdmAeadAlgo::AES_256_GCM,

            SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE,

        );

        #[cfg(feature = "hashed-transcript-data")]

        {

            let mut shared_secret = SpdmSharedSecretFinalKeyStruct::default();

            shared_secret.data_size = SpdmDheAlgo::SECP_384_R1.get_key_size();

            requester.common.session[0]

                .set_shared_secret(SpdmVersion::SpdmVersion12, shared_secret)

                .unwrap();

            requester.common.session[0].runtime_info.digest_context_th =

                spdmlib::crypto::hash::hash_ctx_init(SpdmBaseHashAlgo::TPM_ALG_SHA_384);

        }

        requester.common.session[0].set_session_state(SpdmSessionState::SpdmSessionHandshaking);

        let _ = requester.send_receive_spdm_finish(None, 4294836221).await;

    }

    // TCD:

    // - id: 0

    // - title: 'Fuzz SPDM handle finish response'

    // - description: '<p>Request finish to complete the handshake, and the handshake messages are secured.</p>'

    // -

    {

        let (req_config_info, req_provision_info) = req_create_info();

        let shared_buffer = SharedBuffer::new();

        let pcidoe_transport_encap = Arc::new(Mutex::new(PciDoeTransportEncap {}));

        let mut device_io_requester =

            fake_device_io::FakeSpdmDeviceIo::new(Arc::new(shared_buffer));

        device_io_requester.set_rx(&fuzzdata);

        let device_io_requester = Arc::new(Mutex::new(device_io_requester));

        let mut requester = requester::RequesterContext::new(

            device_io_requester,

            pcidoe_transport_encap,

            req_config_info,

            req_provision_info,

        );

        requester.common.negotiate_info.spdm_version_sel = SpdmVersion::SpdmVersion12;

        requester.common.negotiate_info.req_ct_exponent_sel = 0;

        requester.common.negotiate_info.req_capabilities_sel =

            SpdmRequestCapabilityFlags::CERT_CAP | SpdmRequestCapabilityFlags::KEY_UPD_CAP;

        requester.common.negotiate_info.rsp_ct_exponent_sel = 0;

        requester.common.negotiate_info.rsp_capabilities_sel =

            SpdmResponseCapabilityFlags::CERT_CAP | SpdmResponseCapabilityFlags::KEY_UPD_CAP;

        requester.common.negotiate_info.base_hash_sel = SpdmBaseHashAlgo::TPM_ALG_SHA_384;

        requester.common.negotiate_info.base_asym_sel =

            SpdmBaseAsymAlgo::TPM_ALG_ECDSA_ECC_NIST_P384;

        requester.common.negotiate_info.dhe_sel = SpdmDheAlgo::SECP_384_R1;

        requester.common.negotiate_info.aead_sel = SpdmAeadAlgo::AES_256_GCM;

        requester.common.negotiate_info.req_asym_sel = SpdmReqAsymAlgo::TPM_ALG_RSAPSS_2048;

        requester.common.negotiate_info.key_schedule_sel = SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE;

        requester.common.reset_runtime_info();

        requester.common.session[0] = SpdmSession::new();

        requester.common.session[0].setup(4294836221).unwrap();

        requester.common.session[0].set_crypto_param(

            SpdmBaseHashAlgo::TPM_ALG_SHA_384,

            SpdmDheAlgo::SECP_384_R1,

            SpdmKemAlgo::empty(),

            SpdmAeadAlgo::AES_256_GCM,

            SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE,

        );

        #[cfg(feature = "hashed-transcript-data")]

        {

            let mut shared_secret = SpdmSharedSecretFinalKeyStruct::default();

            shared_secret.data_size = SpdmDheAlgo::SECP_384_R1.get_key_size();

            requester.common.session[0]

                .set_shared_secret(SpdmVersion::SpdmVersion12, shared_secret)

                .unwrap();

            requester.common.session[0].runtime_info.digest_context_th =

                spdmlib::crypto::hash::hash_ctx_init(SpdmBaseHashAlgo::TPM_ALG_SHA_384);

        }

        requester.common.session[0].set_session_state(SpdmSessionState::SpdmSessionHandshaking);

        let _ = requester.send_receive_spdm_finish(None, 4294836221).await;

    }

    // TCD:

    // - id: 0

    // - title: 'Fuzz SPDM handle finish response'

    // - description: '<p>Request finish to complete the handshake with mut auth requested.</p>'

    // -

    {

        let (req_config_info, req_provision_info) = req_create_info();

        let shared_buffer = SharedBuffer::new();

        let pcidoe_transport_encap = Arc::new(Mutex::new(PciDoeTransportEncap {}));

        let mut device_io_requester =

            fake_device_io::FakeSpdmDeviceIo::new(Arc::new(shared_buffer));

        device_io_requester.set_rx(&fuzzdata);

        let device_io_requester = Arc::new(Mutex::new(device_io_requester));

        let mut requester = requester::RequesterContext::new(

            device_io_requester,

            pcidoe_transport_encap,

            req_config_info,

            req_provision_info,

        );

        requester.common.negotiate_info.spdm_version_sel = SpdmVersion::SpdmVersion12;

        requester.common.negotiate_info.req_ct_exponent_sel = 0;

        requester.common.negotiate_info.req_capabilities_sel =

            SpdmRequestCapabilityFlags::CERT_CAP | SpdmRequestCapabilityFlags::KEY_UPD_CAP;

        requester.common.negotiate_info.rsp_ct_exponent_sel = 0;

        requester.common.negotiate_info.rsp_capabilities_sel =

            SpdmResponseCapabilityFlags::CERT_CAP | SpdmResponseCapabilityFlags::KEY_UPD_CAP;

        requester.common.negotiate_info.base_hash_sel = SpdmBaseHashAlgo::TPM_ALG_SHA_384;

        requester.common.negotiate_info.base_asym_sel =

            SpdmBaseAsymAlgo::TPM_ALG_ECDSA_ECC_NIST_P384;

        requester.common.negotiate_info.dhe_sel = SpdmDheAlgo::SECP_384_R1;

        requester.common.negotiate_info.aead_sel = SpdmAeadAlgo::AES_256_GCM;

        requester.common.negotiate_info.req_asym_sel = SpdmReqAsymAlgo::TPM_ALG_ECDSA_ECC_NIST_P384;

        requester.common.negotiate_info.key_schedule_sel = SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE;

        requester.common.peer_info.peer_cert_chain[0] = Some(get_rsp_cert_chain_buff());

        requester.common.provision_info.my_cert_chain = [

            Some(get_rsp_cert_chain_buff()),

            None,

            None,

            None,

            None,

            None,

            None,

            None,

        ];

        requester.common.reset_runtime_info();

        requester.common.session[0] = SpdmSession::new();

        requester.common.session[0].setup(4294836221).unwrap();

        requester.common.session[0].set_crypto_param(

            SpdmBaseHashAlgo::TPM_ALG_SHA_384,

            SpdmDheAlgo::SECP_384_R1,

            SpdmKemAlgo::empty(),

            SpdmAeadAlgo::AES_256_GCM,

            SpdmKeyScheduleAlgo::SPDM_KEY_SCHEDULE,

        );

        #[cfg(feature = "hashed-transcript-data")]

        {

            let mut shared_secret = SpdmSharedSecretFinalKeyStruct::default();

            shared_secret.data_size = SpdmDheAlgo::SECP_384_R1.get_key_size();

            requester.common.session[0]

                .set_shared_secret(SpdmVersion::SpdmVersion12, shared_secret)

                .unwrap();

            requester.common.session[0].runtime_info.digest_context_th =

                spdmlib::crypto::hash::hash_ctx_init(SpdmBaseHashAlgo::TPM_ALG_SHA_384);

        }

        requester.common.session[0].set_session_state(SpdmSessionState::SpdmSessionHandshaking);

        requester.common.session[0].set_mut_auth_requested(

            SpdmKeyExchangeMutAuthAttributes::MUT_AUTH_REQ_WITH_GET_DIGESTS,

        );

        let _ = requester

            .send_receive_spdm_finish(Some(0), 4294836221)

            .await;

    }

}

#[cfg(not(feature = "use_libfuzzer"))]

fn main() {

    #[cfg(all(feature = "fuzzlogfile", feature = "fuzz"))]

    flexi_logger::Logger::try_with_str("info")

        .unwrap()

        .log_to_file(

            FileSpec::default()

                .directory("traces")

                .basename("foo")

                .discriminant("Sample4711A")

                .suffix("trc"),

        )

        .print_message()

        .create_symlink("current_run")

        .start()

        .unwrap();

    #[cfg(not(feature = "fuzz"))]

    {

        let args: Vec<String> = std::env::args().collect();

        if args.len() < 2 {

            // Here you can replace the single-step debugging value in the fuzzdata array.

            let fuzzdata = vec![

                0x1, 0x0, 0x2, 0x0, 0x9, 0x0, 0x0, 0x0, 0xfe, 0xff, 0xfe, 0xff, 0x16, 0x0, 0xca,

                0xa7, 0x51, 0x58, 0x4d, 0x60, 0xe6, 0xc5, 0x74, 0x1c, 0xb3, 0xae, 0xaf, 0x62, 0x4b,

                0x2e, 0x49, 0x54, 0x7a, 0x75, 0x86, 0x37,

            ];

            executor::block_on(fuzz_send_receive_spdm_finish(Arc::new(fuzzdata)));

        } else {

            let path = &args[1];

            let data = std::fs::read(path).expect("read crash file fail");

            executor::block_on(fuzz_send_receive_spdm_finish(Arc::new(data)));

        }

    }

    #[cfg(feature = "fuzz")]

    afl::fuzz!(|data: &[u8]| {

        executor::block_on(fuzz_send_receive_spdm_finish(Arc::new(data.to_vec())));

    });

}