rs-matter 0.3.0

/*
 *
 *    Copyright (c) 2026 Project CHIP Authors
 *
 *    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.
 */

#![cfg(all(feature = "std", feature = "async-io"))]

use core::pin::pin;

use embassy_futures::select::{select, Either};
use embassy_time::{Duration, Timer};

use log::info;
use std::sync::Mutex;

use rs_matter::cert::gen::VALID_FOREVER;
use rs_matter::cert::MAX_CERT_TLV_AND_ASN1_LEN;
use rs_matter::crypto::{
    test_only_crypto, CanonAeadKey, CanonAeadKeyRef, CanonPkcSecretKey, Crypto, RngCore, SecretKey,
    SigningSecretKey, AEAD_CANON_KEY_LEN,
};
use rs_matter::dm::devices::test::{TEST_DEV_ATT, TEST_DEV_COMM, TEST_DEV_DET};
use rs_matter::error::Error;
use rs_matter::onboard::cac::RcacGenerator;
use rs_matter::onboard::noc::NocGenerator;
use rs_matter::respond::Responder;
use rs_matter::sc::case::CaseInitiator;
use rs_matter::sc::{OpCode, SecureChannel, PROTO_ID_SECURE_CHANNEL};
use rs_matter::transport::exchange::Exchange;
use rs_matter::transport::network::{Address, NetworkSend, NoNetwork};
use rs_matter::transport::packet::PacketHdr;

use rs_matter::utils::select::Coalesce;
use rs_matter::utils::storage::ParseBuf;
use rs_matter::Matter;

use crate::common::{create_localhost_socket_pair, init_env_logger, run_device_controller};

#[allow(dead_code)]
mod common;

const TEST_FABRIC_ID: u64 = 1;
const CONTROLLER_NODE_ID: u64 = 100;
const DEVICE_NODE_ID: u64 = 200;

struct DropFirstSigma2<'a> {
    socket: &'a async_io::Async<std::net::UdpSocket>,
    enabled: bool,
    first_two_packets: &'a Mutex<Vec<Vec<u8>>>,
}

impl NetworkSend for DropFirstSigma2<'_> {
    async fn send_to(&mut self, data: &[u8], addr: Address) -> Result<(), Error> {
        if self.enabled {
            let mut packets = self.first_two_packets.lock().unwrap();

            if packets.is_empty() && is_sigma2(data) {
                packets.push(data.to_vec());

                // Simulate loss of the initial Sigma2.
                return Ok(());
            } else if packets.len() == 1 {
                if is_sigma2(data) {
                    packets.push(data.to_vec());
                } else {
                    // Do not let a later standalone ACK advance the initiator's
                    // peer message counter past the pending Sigma2 retry.
                    return Ok(());
                }
            }
        }

        let mut socket = self.socket;
        NetworkSend::send_to(&mut socket, data, addr).await
    }
}

fn is_sigma2(data: &[u8]) -> bool {
    let mut data = data.to_vec();
    let mut pb = ParseBuf::new(data.as_mut_slice());
    let mut header = PacketHdr::new();

    if header.decode_plain_hdr(&mut pb).is_err()
        || header
            .decode_remaining(test_only_crypto(), None, 0, &mut pb)
            .is_err()
    {
        return false;
    }

    header.proto.proto_id == PROTO_ID_SECURE_CHANNEL
        && header.proto.proto_opcode == OpCode::CASESigma2 as u8
}

/// Test that a full CASE handshake succeeds between two in-process Matter instances.
///
/// The controller initiates a CASE session using `CaseInitiator`, while the device
/// runs a `SecureChannel` responder that handles the handshake. On success the
/// unsecured session is upgraded to a secure CASE session.
///
/// Uses [`RcacGenerator`] + [`NocGenerator`] directly to mint a shared
/// RCAC and the controller + device NOCs against it, then plants the
/// resulting (key, RCAC, NOC, IPK) tuple in **both** Matter instances'
/// fabric tables via `Fabrics::add`. Same fabric, two member nodes.
///
/// In-process tests like this one need to fake the AddNOC step on the
/// device side — there is no commissioner here, just two peers that
/// must agree on the same fabric. Real commissioning is exercised by
/// `tests/commissioning.rs`.
#[test]
fn test_case_handshake() {
    run_case_handshake_test(false);
}

/// Test that retransmitting Sigma2 produces the same packet and still completes CASE.
///
/// `Exchange::send_with` rebuilds a packet when MRP requests a retransmission.
/// Dropping the responder's first Sigma2 forces that path.
#[test]
fn test_case_handshake_with_sigma2_retransmission() {
    run_case_handshake_test(true);
}

fn run_case_handshake_test(drop_first_sigma2: bool) {
    init_env_logger();

    futures_lite::future::block_on(async {
        let crypto = test_only_crypto();

        // ---- 1. Generate shared fabric material: RCAC + IPK + NOCs ----

        // Build a shared RCAC (RCAC-direct mode, no ICAC tier) and
        // a single NocGenerator that signs both NOCs against it —
        // controller and device land on the same fabric.
        let mut rcac_buf = [0u8; MAX_CERT_TLV_AND_ASN1_LEN];
        let mut rcac_gen = RcacGenerator::new(&mut rcac_buf);
        let (rcac_privkey, rcac) = rcac_gen
            .generate(&crypto, TEST_FABRIC_ID, VALID_FOREVER)
            .unwrap();

        let mut noc_buf = [0u8; MAX_CERT_TLV_AND_ASN1_LEN];
        let mut noc_generator =
            NocGenerator::create(rcac_privkey.reference(), rcac, &[], &mut noc_buf).unwrap();

        // Shared fabric IPK (16 random bytes).
        let mut ipk = CanonAeadKey::new();
        let mut ipk_bytes = [0u8; AEAD_CANON_KEY_LEN];
        crypto.rand().unwrap().fill_bytes(&mut ipk_bytes);
        ipk.load_from_array(&ipk_bytes);
        let ipk_ref: CanonAeadKeyRef<'_> = ipk.reference();

        // Controller signing keypair + CSR.
        let controller_secret_key = crypto.generate_secret_key().unwrap();
        let mut controller_csr_buf = [0u8; 256];
        let controller_csr = controller_secret_key.csr(&mut controller_csr_buf).unwrap();
        let mut controller_secret_key_canon = CanonPkcSecretKey::new();
        controller_secret_key
            .write_canon(&mut controller_secret_key_canon)
            .unwrap();

        // Device signing keypair + CSR.
        let device_secret_key = crypto.generate_secret_key().unwrap();
        let mut device_csr_buf = [0u8; 256];
        let device_csr = device_secret_key.csr(&mut device_csr_buf).unwrap();
        let mut device_secret_key_canon = CanonPkcSecretKey::new();
        device_secret_key
            .write_canon(&mut device_secret_key_canon)
            .unwrap();

        // ---- 2. Set up two Matter instances ----

        let device_matter = Matter::new(&TEST_DEV_DET, TEST_DEV_COMM, &TEST_DEV_ATT, 0);
        let controller_matter = Matter::new(&TEST_DEV_DET, TEST_DEV_COMM, &TEST_DEV_ATT, 0);

        // ---- 3. Install the same fabric in both fabric tables ----
        //
        // Sign-and-install one side at a time: the NOC slice returned
        // by `noc_generator.generate` borrows `noc_buf`, so it must be
        // consumed by `Fabrics::add` before the next `generate` call
        // overwrites the buffer.

        let controller_noc = noc_generator
            .generate(
                &crypto,
                controller_csr,
                CONTROLLER_NODE_ID,
                &[],
                VALID_FOREVER,
            )
            .unwrap();

        let controller_fab_idx = controller_matter.with_state(|state| {
            state
                .fabrics
                .add(
                    &crypto,
                    controller_secret_key_canon.reference(),
                    rcac,
                    controller_noc,
                    &[], // no ICAC
                    Some(ipk_ref),
                    0xFFF1,
                    CONTROLLER_NODE_ID,
                )
                .unwrap()
                .fab_idx()
        });

        let device_noc = noc_generator
            .generate(&crypto, device_csr, DEVICE_NODE_ID, &[], VALID_FOREVER)
            .unwrap();

        device_matter.with_state(|state| {
            state
                .fabrics
                .add(
                    &crypto,
                    device_secret_key_canon.reference(),
                    rcac,
                    device_noc,
                    &[], // no ICAC
                    Some(ipk_ref),
                    0xFFF1,
                    CONTROLLER_NODE_ID,
                )
                .unwrap();
        });

        // ---- 4. Bind UDP sockets ----

        let (device_socket, controller_socket) = create_localhost_socket_pair();
        let peer_addr = Address::Udp(device_socket.get_ref().local_addr().unwrap());

        // ---- 5. Device side: transport + SecureChannel responder ----
        // No need to open commissioning window for CASE

        let sc = SecureChannel::new(&crypto, &());
        let responder = Responder::new("device", sc, &device_matter, 0);
        let first_two_packets = Mutex::new(Vec::new());

        let device_fut = async {
            select(
                device_matter.run(
                    &crypto,
                    DropFirstSigma2 {
                        socket: &device_socket,
                        enabled: drop_first_sigma2,
                        first_two_packets: &first_two_packets,
                    },
                    &device_socket,
                    NoNetwork,
                ),
                responder.run::<4>(),
            )
            .coalesce()
            .await
        };

        // ---- 6. Controller side: transport + CASE handshake ----

        let controller_fut = async {
            let mut transport = pin!(controller_matter.run(
                &crypto,
                &controller_socket,
                &controller_socket,
                NoNetwork,
            ));
            let mut test = pin!(run_case_handshake(
                &controller_matter,
                &crypto,
                peer_addr,
                controller_fab_idx,
                DEVICE_NODE_ID,
            ));

            match select(&mut transport, &mut test).await {
                Either::First(transport_result) => {
                    panic!("Controller transport exited prematurely: {transport_result:?}");
                }
                Either::Second(test_result) => {
                    // Give transport a moment to flush final messages
                    let mut flush = pin!(Timer::after(Duration::from_millis(300)));
                    if let Either::First(transport_result) =
                        select(&mut transport, &mut flush).await
                    {
                        panic!("Controller transport error during flush: {transport_result:?}");
                    }
                    test_result
                }
            }
        };

        // ---- 7. Run device and controller concurrently ----

        let result = run_device_controller(device_fut, controller_fut).await;

        if drop_first_sigma2 {
            let packets = first_two_packets.lock().unwrap();

            assert_eq!(
                packets.len(),
                2,
                "expected the initial Sigma2 and one retransmission"
            );
            assert_eq!(
                packets[0], packets[1],
                "CASE Sigma2 must be byte-identical when retransmitted"
            );
        }

        result.unwrap();
    });
}

async fn run_case_handshake<C: Crypto>(
    matter: &Matter<'_>,
    crypto: &C,
    peer_addr: Address,
    fab_idx: core::num::NonZeroU8,
    peer_node_id: u64,
) -> Result<(), Error> {
    info!("Creating plaintext session and initiating CASE handshake...");

    let exchange = Exchange::initiate_plaintext(matter, crypto, peer_addr).await?;
    info!("Exchange initiated: {}", exchange.id());

    info!("Starting CASE handshake...");

    let mut case_fut = pin!(CaseInitiator::perform(
        exchange,
        crypto,
        fab_idx,
        peer_node_id,
    ));
    let mut timeout = pin!(Timer::after(Duration::from_secs(30)));

    let result = match select(&mut case_fut, &mut timeout).await {
        Either::First(result) => result,
        Either::Second(_) => panic!("CASE handshake timed out after 30 seconds"),
    };

    result?;

    info!("CASE handshake completed successfully - secure session established");
    Ok(())
}