lib/src/tests.rs

/*
 * Copyright (C) 2025 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Unit tests for AuthMgr BE
use crate::server::{AuthMgrMessage, AUTHMGR_SERVICE_PORT};
use alloc::vec::Vec;
use android_hardware_security_see_authmgr::aidl::android::hardware::security::see::authmgr::{
    DiceChainEntry::DiceChainEntry as AidlDiceChainEntry, DiceLeafArtifacts::DiceLeafArtifacts,
    DicePolicy::DicePolicy, Error::Error, ExplicitKeyDiceCertChain::ExplicitKeyDiceCertChain,
    IAuthMgrAuthorization::IAuthMgrAuthorization, SignedConnectionRequest::SignedConnectionRequest,
};
use authgraph_boringssl::{BoringEcDsa, BoringRng};
use authgraph_core::key::{CertChain, DiceChainEntry};
use authgraph_core::traits::Rng;
use authgraph_core_test::{
    create_dice_cert_chain_for_guest_os, create_dice_leaf_cert, SAMPLE_INSTANCE_HASH,
};
use authmgr_common::{
    signed_connection_request::{
        ConnectionRequest, TEMP_AUTHMGR_BE_TRANSPORT_ID, TEMP_AUTHMGR_FE_TRANSPORT_ID,
    },
    CMD_RAW, CMD_RPC, TOKEN_LENGTH,
};
use authmgr_common_util::{
    get_constraint_spec_for_static_trusty_ta, get_constraints_spec_for_trusty_vm,
    policy_for_dice_node,
};
use binder::Strong;
use coset::CborSerializable;
use hello_world_trusted_aidl::aidl::android::trusty::trustedhal::IHelloWorld::IHelloWorld;
use rpcbinder::RpcSession;
use test::assert_ok;
use tipc::{Deserialize, Handle, Serialize, Serializer, TipcError};
use trusty_std::alloc::TryAllocFrom;

test::init!();

#[test]
fn test_authmgr_connection() {
    assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
}

#[test]
fn test_authmgr_rpc_command() {
    // Test the AuthMgr Init Service by sending just CMD_RPC over the connection, expect success.

    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
}

#[test]
fn test_authmgr_rpc_command_with_session_setup() {
    // Test the AuthMgr Init Service by sending CMD_RPC over the connection and then setting
    // up session, expect success.

    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let _rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
}

#[test]
fn test_authmgr_init_auth_ok() {
    // Test the IAuthMgrAuthorization AIDL interface - initAuthentication, expect success.
    // Connect and send message indicating the intent to connect to the RPC service

    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (_signing_key, _cdi_values, cert_chain) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain }, None);
    let _challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");
}

#[test]
fn test_authmgr_init_auth_with_invalid_dice_chain() {
    // Test the IAuthMgrAuthorization AIDL interface - initAuthentication with an invalid
    // DICE chain, expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_init_auth = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: Vec::<u8>::new() }, None);
    assert!(result_init_auth.is_err());
    assert_eq!(
        result_init_auth.err().unwrap().service_specific_error(),
        Error::INVALID_DICE_CERT_CHAIN.0
    );
}

#[test]
fn test_authmgr_init_auth_no_instance_id() {
    // Test the IAuthMgrAuthorization AIDL interface - initAuthentication with no instance id,
    // expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) without instance hash
    let (_signing_key, _cdi_values, cert_chain) = create_dice_cert_chain_for_guest_os(None, 1);
    let result_init_auth = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain }, None);
    // Expect error because the instance hash is neither in the DICE chain nor provided externally
    assert!(result_init_auth.is_err());
    assert_eq!(
        result_init_auth.err().unwrap().service_specific_error(),
        Error::INVALID_INSTANCE_IDENTIFIER.0
    );
}

#[test]
fn test_authmgr_duplicate_init_auth_with_same_vm_id() {
    // Test the IAuthMgrAuthorization AIDL interface - with two calls to initAuthentication, from
    // the same transport id (i.e. VM-ID), expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (_signing_key, _cdi_values, cert_chain) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain.clone() }, None);
    let _challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let result_init_auth2 = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain }, None);
    assert!(result_init_auth2.is_err());
    assert_eq!(
        result_init_auth2.err().unwrap().service_specific_error(),
        Error::AUTHENTICATION_ALREADY_STARTED.0
    );
}

#[test]
fn test_authmgr_duplicate_init_auth_with_same_vm_id_after_cache_cleanup() {
    // Test the IAuthMgrAuthorization AIDL interface - with two calls to initAuthentication, from
    // the same transport id (i.e. VM-ID), but after cache cleanup triggered by closing the initial
    // connection, expect success.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (_signing_key, _cdi_values, cert_chain) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain.clone() }, None);
    let _challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");
    // Drop the first connection to trigger cache cleanup
    core::mem::drop(conn_rpc);

    // Setup a new connection
    let conn_rpc_2 =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    assert_ok!(conn_rpc_2.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb_authmgr_2 = || {
        let fd = conn_rpc_2.as_raw_fd();
        Some(fd)
    };

    // Setup RPC connection to the AuthMgr service and execute step 1 of phase 1
    let rpc_session_2 = RpcSession::new();
    let rpc_session_2: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session_2.setup_preconnected_client(cb_authmgr_2),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_init_auth2 = rpc_session_2
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain }, None);
    let _challenge: [u8; TOKEN_LENGTH] = assert_ok!(result_init_auth2);
}

#[test]
fn test_authmgr_complete_auth_ok() {
    // Test the IAuthMgrAuthorization AIDL interface - completeAuthentication, expect success.
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);
}

#[test]
fn test_authmgr_duplicate_init_auth_on_authenticated_connection() {
    // Test the IAuthMgrAuthorization AIDL interface - with two calls to initAuthentication, over
    // an already authenticated connection, expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // ******* Invoke step 1 of phase 1 of the protocol over the same connection ********
    let result_init_auth2 = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes }, None);
    assert!(result_init_auth2.is_err());
    assert_eq!(
        result_init_auth2.err().unwrap().service_specific_error(),
        Error::INSTANCE_ALREADY_AUTHENTICATED.0
    );
}

#[test]
fn test_authmgr_duplicate_init_auth_with_same_instance_id_of_authenticatd_vm_on_new_connection() {
    // Test the IAuthMgrAuthorization AIDL interface - with two calls to initAuthentication, using
    // the same instance id of an authenticated pvm, over a new connection, expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // ******* Invoke step 1 of phase 1 of the protocol over a different connection ********
    let conn_rpc_2 = assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT));
    assert_ok!(conn_rpc_2.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb_authmgr_2 = || {
        let fd = conn_rpc_2.as_raw_fd();
        Some(fd)
    };

    // Setup RPC connection to the AuthMgr service and execute step 1 of phase 1
    let rpc_session_2 = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session_2.setup_preconnected_client(cb_authmgr_2),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_init_auth2 = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes }, None);
    assert!(result_init_auth2.is_err());
    assert_eq!(
        result_init_auth2.err().unwrap().service_specific_error(),
        Error::INSTANCE_ALREADY_AUTHENTICATED.0
    );
}

#[test]
fn test_authmgr_duplicate_init_auth_with_diff_instance_ids_same_vm_id_of_authenticatd_vm_on_new_connection(
) {
    // Test the IAuthMgrAuthorization AIDL interface - with two calls to initAuthentication, using
    // different instance ids, but the same vm-id of an authenticated pvm, over a new connection,
    // expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // Create a DICE chain with a different instance hash
    pub const DIFF_INSTANCE_HASH: [u8; 64] = [
        0x5b, 0x3f, 0xc9, 0x6b, 0xe3, 0x95, 0x59, 0x40, 0x21, 0x09, 0x9c, 0xf3, 0xcd, 0xc7, 0xa4,
        0x2a, 0x7d, 0x7e, 0xf5, 0x8e, 0xd6, 0x4d, 0x82, 0x25, 0x1a, 0x51, 0x27, 0x9d, 0x55, 0x8a,
        0xe9, 0x90, 0xf5, 0x8e, 0xd6, 0x4d, 0x84, 0x25, 0x1a, 0x51, 0x27, 0x9d, 0x5b, 0x3f, 0xc9,
        0x6a, 0xe3, 0x95, 0x59, 0x40, 0x21, 0x09, 0x3d, 0xf3, 0xcd, 0xc7, 0xa4, 0x2a, 0x7d, 0x7e,
        0xf5, 0x8e, 0xf5, 0x8e,
    ];
    let (_signing_key_2, _cdi_values_2, cert_chain_bytes_2) =
        create_dice_cert_chain_for_guest_os(Some(DIFF_INSTANCE_HASH), 1);

    // ******* Invoke step 1 of phase 1 of the protocol over a different connection ********
    let conn_rpc_2 = assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT));
    assert_ok!(conn_rpc_2.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb_authmgr_2 = || {
        let fd = conn_rpc_2.as_raw_fd();
        Some(fd)
    };

    // Setup RPC connection to the AuthMgr service and execute step 1 of phase 1
    let rpc_session_2 = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session_2.setup_preconnected_client(cb_authmgr_2),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_init_auth2 = rpc_session
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes_2 }, None);
    assert!(result_init_auth2.is_err());
    assert_eq!(
        result_init_auth2.err().unwrap().service_specific_error(),
        Error::INSTANCE_ALREADY_AUTHENTICATED.0
    );
}

#[test]
fn test_authmgr_raw_command_without_authentication() {
    // Test the AuthMgr Init Service by sending CMD_RAW over the connection, which should result in
    // an error and connection close at the service side because the previous steps of the protocol
    // have not been performed yet. The only way we can assert that this results in an error in the
    // service is by waiting for `recv` on the connection and trying to send another command and
    // asserting that it fails.
    let conn_raw =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");

    let mut token = [0u8; TOKEN_LENGTH];
    let boring_rng = BoringRng;
    boring_rng.fill_bytes(&mut token);
    let mut msg = Vec::<u8>::new();
    msg.push(CMD_RAW);
    msg.extend_from_slice(&token);
    let cmd_raw = AuthMgrMessage(msg);
    assert_ok!(conn_raw.send(&cmd_raw));

    pub struct AuthMgrMessage(pub Vec<u8>);

    impl Deserialize for AuthMgrMessage {
        type Error = TipcError;
        const MAX_SERIALIZED_SIZE: usize = 4000;

        fn deserialize(bytes: &[u8], _handles: &mut [Option<Handle>]) -> Result<Self, TipcError> {
            Ok(AuthMgrMessage(Vec::try_alloc_from(bytes)?))
        }
    }

    impl<'s> Serialize<'s> for AuthMgrMessage {
        fn serialize<'a: 's, S: Serializer<'s>>(
            &'a self,
            serializer: &mut S,
        ) -> Result<S::Ok, S::Error> {
            serializer.serialize_bytes(self.0.as_slice())
        }
    }

    let mut buf = [0u8; 1];
    let _resul: Result<AuthMgrMessage, TipcError> = conn_raw.recv(&mut buf);

    let result2 = conn_raw.send(&cmd_raw);
    assert!(result2.is_err());
}

#[test]
fn test_authmgr_complete_auth_without_init_auth() {
    // Test IAuthMgrAuthorization AIDL interface - by invoking completeAuthentication without
    // invoking initAuthentication before - expect error
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: Vec::new() },
        &DicePolicy { dicePolicy: Vec::new() },
    );
    assert!(result_complete_auth.is_err());
    assert_eq!(
        result_complete_auth.err().unwrap().service_specific_error(),
        Error::AUTHENTICATION_NOT_STARTED.0
    );
}

#[test]
fn test_authmgr_duplicate_complete_auth_on_the_same_connection() {
    // Test IAuthMgrAuthorization AIDL interface - by invoking completeAuthentication on an alerady
    // authenticated connection.
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature.clone() },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.clone().to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    let result_complete_auth2 = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert!(result_complete_auth2.is_err());
    assert_eq!(
        result_complete_auth2.err().unwrap().service_specific_error(),
        Error::INSTANCE_ALREADY_AUTHENTICATED.0
    );
}

#[test]
fn test_authmgr_duplicate_complete_auth_on_new_connection() {
    // Test the IAuthMgrAuthorization AIDL interface - by invoking completeAuthentication from
    // an authenticated pvm, over a new connection, expect error.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature.clone() },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.clone().to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // ******* Invoke step 2 of phase 1 of the protocol over a different connection ********
    let conn_rpc_2 = assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT));
    assert_ok!(conn_rpc_2.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb_authmgr_2 = || {
        let fd = conn_rpc_2.as_raw_fd();
        Some(fd)
    };

    // Setup RPC connection to the AuthMgr service and execute step 1 of phase 1
    let rpc_session_2 = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session_2.setup_preconnected_client(cb_authmgr_2),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_complete_auth_2 = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert!(result_complete_auth_2.is_err());
    assert_eq!(
        result_complete_auth_2.err().unwrap().service_specific_error(),
        Error::INSTANCE_ALREADY_AUTHENTICATED.0
    );
}

#[test]
fn test_authmgr_duplicate_complete_auth_after_cache_cleanup() {
    // Test the IAuthMgrAuthorization AIDL interface - with two attempts for phase 1, from
    // the same transport id (i.e. VM-ID), but after cache cleanup triggered by closing the initial
    // connection, expect success.
    // Connect and send message indicating the intent to connect to the RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, _cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.clone().to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // Drop the first connection to trigger cache cleanup
    core::mem::drop(conn_rpc);

    // Setup a new connection
    let conn_rpc_2 =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    assert_ok!(conn_rpc_2.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb_authmgr_2 = || {
        let fd = conn_rpc_2.as_raw_fd();
        Some(fd)
    };

    // Setup RPC connection to the AuthMgr service and execute step 1 of phase 1
    let rpc_session_2 = RpcSession::new();
    let rpc_session_2: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session_2.setup_preconnected_client(cb_authmgr_2),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );

    let result_init_auth_2 = rpc_session_2
        .initAuthentication(&ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes }, None);
    let challenge_2: [u8; TOKEN_LENGTH] = assert_ok!(result_init_auth_2);

    // Build the connection request to be signed
    let conn_req_2 = ConnectionRequest::new_for_ffa_transport(
        challenge_2,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    let signature_2 = assert_ok!(
        conn_req_2.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );

    let result_complete_auth_2 = rpc_session_2.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature_2 },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth_2);
}

#[test]
fn authmgr_full_protocol_happy_path() {
    // Connect to the TA and send message indicating the intent to connect to the binder RPC service
    let conn_rpc =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to connect to AuthMgr BE.");
    let cmd_rpc = AuthMgrMessage(vec![CMD_RPC]);
    assert_ok!(conn_rpc.send(&cmd_rpc), "Failed to send the command requesting RPC service.");
    let cb = || {
        let fd = conn_rpc.as_raw_fd();
        Some(fd)
    };
    let rpc_session = RpcSession::new();
    let rpc_session: Strong<dyn IAuthMgrAuthorization> = assert_ok!(
        rpc_session.setup_preconnected_client(cb),
        "Failed to setup pre-connected client for the authmgr rpc service."
    );
    // Create a test DICE chain (with CDI secrets for signing) with an instance hash in vm_entry
    let (signing_key, cdi_values, cert_chain_bytes) =
        create_dice_cert_chain_for_guest_os(Some(SAMPLE_INSTANCE_HASH), 1);
    let result_init_auth = rpc_session.initAuthentication(
        &ExplicitKeyDiceCertChain { diceCertChain: cert_chain_bytes.clone() },
        None,
    );
    let challenge: [u8; TOKEN_LENGTH] =
        assert_ok!(result_init_auth, "Failed to invoke initAuthentication.");

    let cert_chain =
        assert_ok!(CertChain::from_slice(&cert_chain_bytes), "Failed to decode the cert chain");
    // Build the connection request to be signed
    let conn_req = ConnectionRequest::new_for_ffa_transport(
        challenge,
        TEMP_AUTHMGR_FE_TRANSPORT_ID,
        TEMP_AUTHMGR_BE_TRANSPORT_ID,
    );
    // Sign the connection request with the DICE CDI secrets
    let ecdsa = BoringEcDsa;
    let verify_key = assert_ok!(cert_chain.validate(&ecdsa), "Failed to validate the cert chain");
    let signing_algorithm = verify_key.get_cose_sign_algorithm();
    let signature = assert_ok!(
        conn_req.sign(&signing_key, &ecdsa, signing_algorithm),
        "Failed to sign connection request"
    );
    // Create a DICE policy
    let constraint_spec = get_constraints_spec_for_trusty_vm();
    let policy = assert_ok!(
        dice_policy_builder::policy_for_dice_chain(&cert_chain_bytes, constraint_spec),
        "Failed to building policy for pvm"
    );

    // ****** Invoke step 2 of phase 1 of the protocol ******
    let result_complete_auth = rpc_session.completeAuthentication(
        &SignedConnectionRequest { signedConnectionRequest: signature },
        &DicePolicy {
            dicePolicy: assert_ok!(policy.to_vec(), "Failed to encode DICE policy for pvm"),
        },
    );
    assert_ok!(result_complete_auth);

    // Connect to the TA and send a message indicating the intent establish a raw connection and
    // send a token
    let conn_raw =
        assert_ok!(Handle::connect(AUTHMGR_SERVICE_PORT), "Failed to setup a raw connection");
    let mut token = [0u8; TOKEN_LENGTH];
    let boring_rng = BoringRng;
    boring_rng.fill_bytes(&mut token);
    let mut msg = Vec::<u8>::new();
    msg.push(CMD_RAW);
    msg.extend_from_slice(&token);
    let cmd_raw = AuthMgrMessage(msg);
    assert_ok!(conn_raw.send(&cmd_raw));

    // ****** Execute phase 2 of the AuthMgr protocol ******
    // Create DICE certificate and a DICE policy for the client
    let leaf_cert_bytes = create_dice_leaf_cert(cdi_values, "keymint", 1);
    let client_constraint_spec_km = get_constraint_spec_for_static_trusty_ta();
    let leaf_cert =
        assert_ok!(DiceChainEntry::from_slice(&leaf_cert_bytes), "Failed to decode leaf cert");
    let client_policy = assert_ok!(
        policy_for_dice_node(&leaf_cert, client_constraint_spec_km),
        "Failed to create policy for leaf cert"
    );
    let result_client_authz = rpc_session.authorizeAndConnectClientToTrustedService(
        &[],
        "HelloService",
        &token,
        &DiceLeafArtifacts {
            diceLeaf: AidlDiceChainEntry { diceChainEntry: leaf_cert_bytes },
            diceLeafPolicy: DicePolicy {
                dicePolicy: assert_ok!(
                    client_policy.to_vec(),
                    "Failed to encode client dice policy."
                ),
            },
        },
    );
    assert_ok!(result_client_authz);

    // Simulate the client connecting to the IHelloWorld example trusted service over the raw
    // connection authorized via AuthMgr protocol execution above
    let cb_trusted_hal = || {
        let fd = conn_raw.as_raw_fd();
        Some(fd)
    };

    let trusted_service_rpc_session = RpcSession::new();
    let trusted_service_rpc_session: Strong<dyn IHelloWorld> = assert_ok!(
        trusted_service_rpc_session.setup_preconnected_client(cb_trusted_hal),
        "Failed to setup pre-connected client for the trusted service."
    );
    let result = assert_ok!(trusted_service_rpc_session.sayHello("Test."));
    assert_eq!("Hello Test.", result);
}