rsipstack 0.5.22

//! Tests for ServerInvite transaction cleanup behavior when dropped in
//! Completed/Confirmed state (RFC 3261 Section 17.2.1).
//!
//! When a ServerInvite transaction sends a 3xx-6xx response and the TU drops
//! the transaction handle, the transaction layer must still:
//! - Auto-reply to INVITE retransmissions (via finished_transactions)
//! - Absorb ACK and clean up waiting_ack

use crate::sip::headers::*;
use crate::sip::{Method, StatusCode, Version};
use crate::transaction::key::{TransactionKey, TransactionRole};
use crate::transaction::transaction::Transaction;
use crate::transaction::TransactionType;
use crate::transport::SipConnection;
use crate::{
    transport::{udp::UdpConnection, TransportLayer},
    EndpointBuilder,
};
use std::time::Duration;
use tokio::time::{sleep, timeout};
use tokio_util::sync::CancellationToken;

fn make_invite(branch: &str, to_tag: &str) -> crate::sip::Request {
    crate::sip::Request {
        method: Method::Invite,
        uri: crate::sip::Uri::try_from("sip:bob@example.com").unwrap(),
        headers: vec![
            Via::new(&format!(
                "SIP/2.0/UDP 127.0.0.1:5060;branch={}",
                branch
            ))
            .into(),
            CSeq::new("1 INVITE").into(),
            From::new("Alice <sip:alice@example.com>;tag=aliceTag123").into(),
            To::new(&format!("Bob <sip:bob@example.com>;tag={}", to_tag)).into(),
            CallId::new("drop-test-call-id@example.com").into(),
            MaxForwards::new("70").into(),
        ]
        .into(),
        version: Version::V2,
        body: Default::default(),
    }
}

fn make_ack_for_invite(invite: &crate::sip::Request) -> crate::sip::Request {
    let mut headers = invite.headers.clone();
    for header in headers.iter_mut() {
        if let crate::sip::Header::CSeq(cseq) = header {
            *cseq = crate::sip::headers::CSeq::new("1 ACK").into();
        }
    }
    crate::sip::Request {
        method: Method::Ack,
        uri: invite.uri.clone(),
        headers,
        version: invite.version.clone(),
        body: Default::default(),
    }
}

/// Unit test: when a ServerInvite transaction is dropped in Completed state,
/// the last_response must be registered in finished_transactions and
/// waiting_ack must NOT be removed.
#[tokio::test]
async fn test_cleanup_server_invite_completed_registers_finished() -> crate::Result<()> {
    let endpoint = super::create_test_endpoint(Some("127.0.0.1:0")).await?;

    let invite = make_invite("z9hG4bK-drop-unit-1", "serverTagAAA");
    let key = TransactionKey::from_request(&invite, TransactionRole::Server)?;
    let tx = Transaction::new_server(
        key.clone(),
        invite.clone(),
        endpoint.inner.clone(),
        None,
    );
    assert_eq!(tx.transaction_type, TransactionType::ServerInvite);

    // Simulate the Completed transition by setting state + last_response directly.
    // We can't call respond() because there's no connection, but we can test
    // the cleanup() logic by setting the relevant fields.
    drop(tx);

    // After dropping in Trying state (no last_response), nothing should be in finished_transactions
    assert!(
        !endpoint.inner.finished_transactions.contains_key(&key),
        "finished_transactions should be empty when dropped in Trying without a response"
    );

    Ok(())
}

/// Unit test: verify that a ServerInvite dropped in Completed state (with a
/// final response) registers in finished_transactions and keeps waiting_ack.
#[tokio::test]
async fn test_cleanup_server_invite_completed_keeps_waiting_ack() -> crate::Result<()> {
    let endpoint = super::create_test_endpoint(Some("127.0.0.1:0")).await?;

    let invite = make_invite("z9hG4bK-drop-unit-2", "serverTagBBB");
    let key = TransactionKey::from_request(&invite, TransactionRole::Server)?;

    let mut tx = Transaction::new_server(
        key.clone(),
        invite.clone(),
        endpoint.inner.clone(),
        None,
    );

    // Manually set the fields that would be set during respond() + transition(Completed)
    let resp = crate::sip::Response {
        status_code: StatusCode::ServiceUnavailable,
        version: Version::V2,
        headers: invite.headers.clone(),
        body: Default::default(),
    };
    tx.last_response = Some(resp.clone());
    tx.state = crate::transaction::TransactionState::Completed;

    // Manually insert into waiting_ack (as transition(Completed) would do)
    let dialog_id = crate::dialog::DialogId::try_from((&resp, TransactionRole::Server))?;
    endpoint.inner.waiting_ack.insert(dialog_id.clone(), key.clone());

    // Drop the transaction - cleanup() should run
    drop(tx);
    sleep(Duration::from_millis(50)).await;

    // waiting_ack should still be present (NOT removed for Completed state)
    assert!(
        endpoint.inner.waiting_ack.contains_key(&dialog_id),
        "waiting_ack must NOT be removed when ServerInvite is dropped in Completed state"
    );

    // finished_transactions should have the response
    let finished = endpoint.inner.finished_transactions.get(&key);
    assert!(
        finished.is_some(),
        "finished_transactions must contain the response for dropped ServerInvite in Completed"
    );
    if let Some(Some(crate::sip::SipMessage::Response(r))) = finished.map(|v| v.value().clone()) {
        assert_eq!(r.status_code, StatusCode::ServiceUnavailable);
    } else {
        panic!("finished_transactions entry should be a Response with 503");
    }

    Ok(())
}

/// Unit test: verify that a ServerInvite that terminates normally
/// (Terminated state) DOES remove waiting_ack but still registers
/// in finished_transactions for Timer J absorption.
#[tokio::test]
async fn test_cleanup_server_invite_terminated_removes_waiting_ack() -> crate::Result<()> {
    let endpoint = super::create_test_endpoint(Some("127.0.0.1:0")).await?;

    let invite = make_invite("z9hG4bK-drop-unit-3", "serverTagCCC");
    let key = TransactionKey::from_request(&invite, TransactionRole::Server)?;

    let mut tx = Transaction::new_server(
        key.clone(),
        invite.clone(),
        endpoint.inner.clone(),
        None,
    );

    let resp = crate::sip::Response {
        status_code: StatusCode::BusyHere,
        version: Version::V2,
        headers: invite.headers.clone(),
        body: Default::default(),
    };
    tx.last_response = Some(resp.clone());
    // Simulate normal termination (ACK → Confirmed → TimerK → Terminated)
    tx.state = crate::transaction::TransactionState::Terminated;

    let dialog_id = crate::dialog::DialogId::try_from((&resp, TransactionRole::Server))?;
    endpoint.inner.waiting_ack.insert(dialog_id.clone(), key.clone());

    drop(tx);
    sleep(Duration::from_millis(50)).await;

    // waiting_ack SHOULD be removed for Terminated state
    assert!(
        !endpoint.inner.waiting_ack.contains_key(&dialog_id),
        "waiting_ack must be removed when ServerInvite terminates normally"
    );

    // finished_transactions should still have the response (Timer J)
    assert!(
        endpoint.inner.finished_transactions.contains_key(&key),
        "finished_transactions should contain the response for Timer J absorption"
    );

    Ok(())
}

/// Integration test: full UDP flow.
/// 1. Client sends INVITE
/// 2. Server replies 503 and drops tx
/// 3. Client retransmits INVITE → auto-gets 503
/// 4. Client sends ACK → absorbed, waiting_ack cleaned up
#[tokio::test]
async fn test_server_invite_drop_retransmission_and_ack() {
    let token = CancellationToken::new();

    // Server endpoint
    let server_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create server connection");
    let server_conn_sip: SipConnection = server_conn.clone().into();
    let server_addr = server_conn_sip.get_addr().clone();

    let tl = TransportLayer::new(token.child_token());
    tl.add_transport(server_conn_sip.clone());

    let endpoint = EndpointBuilder::new()
        .with_user_agent("rsipstack-test")
        .with_transport_layer(tl)
        .build();

    // Client socket
    let client_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create client connection");
    let client_conn_sip: SipConnection = client_conn.clone().into();

    // Start endpoint serve loop
    let endpoint_inner = endpoint.inner.clone();
    let serve_handle = tokio::spawn(async move {
        let _ = endpoint_inner.serve().await;
    });

    // Build INVITE
    let branch = "z9hG4bK-drop-int-branch1";
    let invite = make_invite(branch, "serverTagInt1");

    // --- Step 1: Send INVITE ---
    client_conn_sip
        .send(invite.clone().into(), Some(&server_addr))
        .await
        .expect("send invite");

    // --- Step 2: Server receives tx, replies 503, drops tx ---
    let mut incoming = endpoint.incoming_transactions().expect("incoming");
    let mut tx = timeout(Duration::from_secs(2), incoming.recv())
        .await
        .expect("timeout waiting for incoming transaction")
        .expect("no incoming transaction");
    assert_eq!(tx.original.method, Method::Invite);

    let tx_key = tx.key.clone();

    // Reply 503
    tx.reply(StatusCode::ServiceUnavailable)
        .await
        .expect("reply 503");

    // Drop the transaction (simulating user dropping after reject)
    drop(tx);
    sleep(Duration::from_millis(100)).await;

    // --- Verify state after drop ---
    assert!(
        endpoint.inner.finished_transactions.contains_key(&tx_key),
        "finished_transactions should contain the 503 response after drop"
    );

    // waiting_ack should still have the entry
    assert_eq!(
        endpoint.inner.waiting_ack.len(),
        1,
        "waiting_ack should still have 1 entry for the dropped ServerInvite"
    );

    // --- Step 3: Client receives the 503 response ---
    let mut buf = vec![0u8; 4096];
    let (len, _) = timeout(Duration::from_secs(2), client_conn.recv_raw(&mut buf))
        .await
        .expect("timeout waiting for 503 response")
        .expect("recv failed");
    let resp_str = String::from_utf8_lossy(&buf[..len]);
    assert!(
        resp_str.contains("503"),
        "first response should be 503, got: {}",
        &resp_str[..resp_str.len().min(200)]
    );

    // --- Step 4: Retransmit INVITE → should auto-get 503 ---
    client_conn_sip
        .send(invite.clone().into(), Some(&server_addr))
        .await
        .expect("retransmit invite");

    let (len, _) = timeout(Duration::from_secs(2), client_conn.recv_raw(&mut buf))
        .await
        .expect("timeout waiting for retransmitted 503 response")
        .expect("recv failed on retransmission");
    let resp_str = String::from_utf8_lossy(&buf[..len]);
    assert!(
        resp_str.contains("503"),
        "retransmitted INVITE should get 503 auto-reply, got: {}",
        &resp_str[..resp_str.len().min(200)]
    );

    // --- Step 5: Send ACK → should be absorbed ---
    let ack = make_ack_for_invite(&invite);
    client_conn_sip
        .send(ack.clone().into(), Some(&server_addr))
        .await
        .expect("send ack");

    sleep(Duration::from_millis(200)).await;

    // waiting_ack should now be cleaned up (ACK absorbed)
    assert_eq!(
        endpoint.inner.waiting_ack.len(),
        0,
        "waiting_ack should be empty after ACK is absorbed"
    );

    // finished_transactions may or may not still have the entry (TimerCleanup runs later)
    // but the key behavior is that ACK was absorbed without error

    // Verify no response is sent back for the ACK (client shouldn't receive anything)
    match timeout(Duration::from_millis(300), client_conn.recv_raw(&mut buf)).await {
        Err(_) => { /* good - no response for ACK, timed out */ }
        Ok(Ok((len, _))) => {
            let resp_str = String::from_utf8_lossy(&buf[..len]);
            panic!("should not receive anything after ACK, got: {}", &resp_str[..resp_str.len().min(200)]);
        }
        Ok(Err(e)) => {
            panic!("unexpected error waiting after ACK: {:?}", e);
        }
    }

    serve_handle.abort();
}

/// Integration test: ServerInvite normal termination (ACK received via dialog/handle)
/// also registers in finished_transactions for Timer J.
#[tokio::test]
async fn test_server_invite_normal_termination_registers_finished() {
    let token = CancellationToken::new();

    let server_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create server connection");
    let server_conn_sip: SipConnection = server_conn.clone().into();
    let server_addr = server_conn_sip.get_addr().clone();

    let tl = TransportLayer::new(token.child_token());
    tl.add_transport(server_conn_sip.clone());

    let endpoint = EndpointBuilder::new()
        .with_user_agent("rsipstack-test")
        .with_transport_layer(tl)
        .build();

    let client_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create client connection");
    let client_conn_sip: SipConnection = client_conn.clone().into();

    let endpoint_inner = endpoint.inner.clone();
    let serve_handle = tokio::spawn(async move {
        let _ = endpoint_inner.serve().await;
    });

    let branch = "z9hG4bK-normal-term-branch";
    let invite = make_invite(branch, "serverTagNormal");

    // Send INVITE
    client_conn_sip
        .send(invite.clone().into(), Some(&server_addr))
        .await
        .expect("send invite");

    let mut incoming = endpoint.incoming_transactions().expect("incoming");

    // Server loop: reply 503, then wait for ACK, then drop tx
    let server_loop = async {
        let mut tx = timeout(Duration::from_secs(2), incoming.recv())
            .await
            .expect("timeout waiting for incoming")
            .expect("no incoming");
        assert_eq!(tx.original.method, Method::Invite);

        let tx_key = tx.key.clone();

        // Reply 503
        tx.reply(StatusCode::ServiceUnavailable)
            .await
            .expect("reply 503");

        // Run receive loop to handle ACK
        while let Some(msg) = tx.receive().await {
            if let crate::sip::SipMessage::Request(req) = msg {
                if req.method == Method::Ack {
                    // ACK received, break
                    break;
                }
            }
        }

        // Now drop tx (simulating handle_invite returning)
        drop(tx);
        sleep(Duration::from_millis(100)).await;

        // Verify finished_transactions has the response (Timer J)
        assert!(
            endpoint.inner.finished_transactions.contains_key(&tx_key),
            "finished_transactions should contain the response after normal termination"
        );

        // waiting_ack should have been cleaned up (ACK was processed via Confirmed path
        // but the drop happens while state is Confirmed before TimerK fires, so
        // cleanup keeps it — then the ACK absorption in on_received_message cleaned it.
        // Actually, the ACK was processed by the transaction's receive() loop,
        // transitioning to Confirmed. Then tx is dropped from Confirmed state.
        // Our fix keeps waiting_ack in Confirmed state, but the TimerCleanup
        // will eventually remove it. This is acceptable.)
    };

    let client_loop = async {
        let mut buf = vec![0u8; 4096];

        // Receive 503
        let (len, _) = timeout(Duration::from_secs(2), client_conn.recv_raw(&mut buf))
            .await
            .expect("timeout waiting for 503")
            .expect("recv failed");
        assert!(String::from_utf8_lossy(&buf[..len]).contains("503"));

        sleep(Duration::from_millis(50)).await;

        // Send ACK
        let ack = make_ack_for_invite(&invite);
        client_conn_sip
            .send(ack.into(), Some(&server_addr))
            .await
            .expect("send ack");
    };

    tokio::join!(server_loop, client_loop);

    serve_handle.abort();
}

/// Verify that finished_transactions properly absorbs late INVITE retransmissions
/// with the correct status code kind.
#[tokio::test]
async fn test_finished_transactions_replies_correct_status() {
    let token = CancellationToken::new();

    let server_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create server connection");
    let server_conn_sip: SipConnection = server_conn.clone().into();
    let server_addr = server_conn_sip.get_addr().clone();

    let tl = TransportLayer::new(token.child_token());
    tl.add_transport(server_conn_sip.clone());

    let endpoint = EndpointBuilder::new()
        .with_user_agent("rsipstack-test")
        .with_transport_layer(tl)
        .build();

    let client_conn =
        UdpConnection::create_connection("127.0.0.1:0".parse().unwrap(), None, None)
            .await
            .expect("create client connection");
    let client_conn_sip: SipConnection = client_conn.clone().into();

    let endpoint_inner = endpoint.inner.clone();
    let serve_handle = tokio::spawn(async move {
        let _ = endpoint_inner.serve().await;
    });

    let branch = "z9hG4bK-status-check-branch";
    let invite = make_invite(branch, "serverTagStatus");

    // Send INVITE
    client_conn_sip
        .send(invite.clone().into(), Some(&server_addr))
        .await
        .expect("send invite");

    let mut incoming = endpoint.incoming_transactions().expect("incoming");
    let mut tx = timeout(Duration::from_secs(2), incoming.recv())
        .await
        .expect("timeout waiting for incoming")
        .expect("no incoming");

    // Reply 486 Busy Here
    tx.reply(StatusCode::BusyHere).await.expect("reply 486");
    drop(tx);
    sleep(Duration::from_millis(100)).await;

    let mut buf = vec![0u8; 4096];

    // Receive 486
    let (len, _) = timeout(Duration::from_secs(2), client_conn.recv_raw(&mut buf))
        .await
        .expect("timeout waiting for 486")
        .expect("recv failed");
    let resp = String::from_utf8_lossy(&buf[..len]);
    assert!(resp.contains("486"), "should receive 486");

    // Retransmit INVITE
    sleep(Duration::from_millis(50)).await;
    client_conn_sip
        .send(invite.clone().into(), Some(&server_addr))
        .await
        .expect("retransmit invite");

    // Receive 486 again
    let (len, _) = timeout(Duration::from_secs(2), client_conn.recv_raw(&mut buf))
        .await
        .expect("timeout waiting for retransmitted 486")
        .expect("recv failed");
    let resp = String::from_utf8_lossy(&buf[..len]);
    assert!(
        resp.contains("486"),
        "retransmitted INVITE should get 486 auto-reply"
    );

    serve_handle.abort();
}