simple-someip 0.5.3

use crate::{
    e2e::{E2ECheckStatus, E2EKey, E2ERegistry, PROFILE4_HEADER_SIZE},
    protocol::{Message, MessageView, sd},
    traits::{PayloadWireFormat, WireFormat},
};

use super::error::Error;
use std::{
    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
    sync::{Arc, Mutex},
    task::{Context, Poll},
    vec,
};
use tokio::{net::UdpSocket, select, sync::mpsc};
use tracing::{error, info, trace};

/// A received message together with the source address it came from.
#[derive(Clone, Debug)]
pub struct ReceivedMessage<P> {
    pub message: Message<P>,
    pub source: SocketAddr,
    pub e2e_status: Option<E2ECheckStatus>,
}

/// Structure representing a request to send a message
#[derive(Debug)]
pub struct SendMessage<PayloadDefinitions> {
    pub target_addr: SocketAddrV4,
    pub message: Message<PayloadDefinitions>,
    response: tokio::sync::oneshot::Sender<Result<(), Error>>,
}

impl<PayloadDefinitions: PayloadWireFormat + 'static> SendMessage<PayloadDefinitions> {
    pub fn new(
        target_addr: SocketAddrV4,
        message: Message<PayloadDefinitions>,
    ) -> (tokio::sync::oneshot::Receiver<Result<(), Error>>, Self) {
        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
        (
            response_rx,
            Self {
                target_addr,
                message,
                response: response_tx,
            },
        )
    }
}

#[derive(Debug)]
pub struct SocketManager<PayloadDefinitions> {
    receiver: mpsc::Receiver<Result<ReceivedMessage<PayloadDefinitions>, Error>>,
    sender: mpsc::Sender<SendMessage<PayloadDefinitions>>,
    local_port: u16,
    session_id: u16,
}

impl<MessageDefinitions> SocketManager<MessageDefinitions>
where
    MessageDefinitions: PayloadWireFormat + 'static,
{
    pub fn bind_discovery(
        interface: Ipv4Addr,
        e2e_registry: Arc<Mutex<E2ERegistry>>,
    ) -> Result<Self, Error> {
        let (rx_tx, rx_rx) = mpsc::channel(16);
        let (tx_tx, tx_rx) = mpsc::channel(16);
        let bind_addr =
            std::net::SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), sd::MULTICAST_PORT);

        // Create socket with SO_REUSEADDR to allow quick restart
        let socket = socket2::Socket::new(
            socket2::Domain::IPV4,
            socket2::Type::DGRAM,
            Some(socket2::Protocol::UDP),
        )?;
        socket.set_reuse_address(true)?;
        #[cfg(unix)]
        socket.set_reuse_port(true)?;
        socket.set_multicast_if_v4(&interface)?;
        // Disable multicast loopback so this socket does not receive the
        // SD messages it sends. Matches the Server's SD socket setup —
        // otherwise a client that acts as both server and client (e.g.
        // offering services via `start_sd_announcements`) will parse its
        // own OfferService entries as peer offers.
        socket.set_multicast_loop_v4(false)?;
        socket.bind(&bind_addr.into())?;
        socket.set_nonblocking(true)?;
        let socket: std::net::UdpSocket = socket.into();
        let socket = UdpSocket::from_std(socket)?;

        socket.join_multicast_v4(sd::MULTICAST_IP, interface)?;

        Self::spawn_socket_loop(socket, rx_tx, tx_rx, e2e_registry);
        Ok(Self {
            receiver: rx_rx,
            sender: tx_tx,
            local_port: sd::MULTICAST_PORT,
            session_id: 0,
        })
    }

    pub fn bind(port: u16, e2e_registry: Arc<Mutex<E2ERegistry>>) -> Result<Self, Error> {
        let (rx_tx, rx_rx) = mpsc::channel(4);
        let (tx_tx, tx_rx) = mpsc::channel(4);
        let bind_addr = std::net::SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port);

        // Create socket with SO_REUSEADDR and SO_REUSEPORT to allow quick restart
        let socket = socket2::Socket::new(
            socket2::Domain::IPV4,
            socket2::Type::DGRAM,
            Some(socket2::Protocol::UDP),
        )?;
        socket.set_reuse_address(true)?;
        socket.bind(&bind_addr.into())?;
        socket.set_nonblocking(true)?;
        let socket: std::net::UdpSocket = socket.into();
        let socket = UdpSocket::from_std(socket)?;
        let port = socket.local_addr()?.port();
        Self::spawn_socket_loop(socket, rx_tx, tx_rx, e2e_registry);
        Ok(Self {
            receiver: rx_rx,
            sender: tx_tx,
            local_port: port,
            session_id: 0,
        })
    }

    pub async fn send(
        &mut self,
        target_addr: SocketAddrV4,
        message: Message<MessageDefinitions>,
    ) -> Result<(), Error> {
        let (result_channel, message) = SendMessage::new(target_addr, message);
        self.sender.send(message).await.map_err(|e| {
            error!("Socket error: {e} when attempting to send message");
            Error::SocketClosedUnexpectedly
        })?;
        result_channel
            .await
            .expect("Socket manager must always return result of send before dropping channel")?;
        self.session_id += 1;
        Ok(())
    }

    pub async fn receive(&mut self) -> Option<Result<ReceivedMessage<MessageDefinitions>, Error>> {
        self.receiver.recv().await
    }

    /// Poll the receiver for a message without blocking.
    /// Used by `Inner::receive_any_unicast` to poll multiple sockets.
    pub fn poll_receive(
        &mut self,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<ReceivedMessage<MessageDefinitions>, Error>>> {
        self.receiver.poll_recv(cx)
    }

    pub fn session_id(&self) -> u16 {
        self.session_id
    }

    pub fn port(&self) -> u16 {
        self.local_port
    }

    pub async fn shut_down(self) {
        let Self {
            sender,
            mut receiver,
            ..
        } = self;
        drop(sender);
        _ = receiver.recv().await;
    }

    #[allow(clippy::too_many_lines)]
    fn spawn_socket_loop(
        socket: UdpSocket,
        rx_tx: mpsc::Sender<Result<ReceivedMessage<MessageDefinitions>, Error>>,
        mut tx_rx: mpsc::Receiver<SendMessage<MessageDefinitions>>,
        e2e_registry: Arc<Mutex<E2ERegistry>>,
    ) {
        tokio::spawn(async move {
            let mut buf = vec![0; 1400];
            loop {
                select! {
                    result = socket.recv_from(&mut buf) => {
                        match result {
                            Ok((bytes_received, source_address)) => {
                                let parse_result = MessageView::parse(&buf[..bytes_received])
                                    .and_then(|view| {
                                        let header = view.header().to_owned();
                                        let upper_header = header.upper_header_bytes();
                                        let key = E2EKey::from_message_id(header.message_id());
                                        let payload_bytes = view.payload_bytes();

                                        // Apply E2E check if configured
                                        let (e2e_status, effective_payload) = {
                                            let mut registry = e2e_registry.lock().expect("e2e registry lock poisoned");
                                            match registry.check(key, payload_bytes, upper_header) {
                                                Some((status, stripped)) => (Some(status), stripped),
                                                None => (None, payload_bytes),
                                            }
                                        };

                                        let payload = MessageDefinitions::from_payload_bytes(header.message_id(), effective_payload)?;
                                        Ok(ReceivedMessage {
                                            message: Message::new(header, payload),
                                            source: source_address,
                                            e2e_status,
                                        })
                                    })
                                    .map_err(Error::from);
                                if let Ok(()) = rx_tx.send( parse_result ).await {} else {
                                    info!("Socket Dropping");
                                    // The receiver has been dropped, so we should exit
                                    break;
                                }
                            }
                            Err(e) => {

                                error!("Error decoding message: {:?}", e);
                            }
                        }
                    },
                    message = tx_rx.recv() => {
                        if let Some(send_message) = message {
                            trace!("Sending: {:?}", &send_message);
                            let mut message_length = match send_message.message.encode(&mut buf.as_mut_slice()) {
                                Ok(length) => length,
                                Err(e) => {
                                    error!("Failed to encode message: {:?}", e);
                                    // If the sender is already closed we can't send the error back, so we shut everything down
                                    if let Ok(()) = send_message.response.send(Err(e.into())) {
                                        // Successfully sent error back to sender, carry on
                                        continue;
                                    }
                                    error!("Socket owner closed channel unexpectedly, closing socket.");
                                    break;
                                }
                            };

                            // Apply E2E protect if configured
                            {
                                let key = E2EKey::from_message_id(send_message.message.header().message_id());
                                let mut registry = e2e_registry.lock().expect("e2e registry lock poisoned");
                                if registry.contains_key(&key) {
                                    let original_payload = buf[16..message_length].to_vec();
                                    let upper_header: [u8; 8] = buf[8..16].try_into().expect("upper header slice");
                                    let mut protected = vec![0u8; original_payload.len() + PROFILE4_HEADER_SIZE];
                                    match registry.protect(key, &original_payload, upper_header, &mut protected) {
                                        Some(Ok(protected_len)) => {
                                            #[allow(clippy::cast_possible_truncation)]
                                            let new_length: u32 = 8 + protected_len as u32;
                                            buf[4..8].copy_from_slice(&new_length.to_be_bytes());
                                            if 16 + protected_len > buf.len() {
                                                buf.resize(16 + protected_len, 0);
                                            }
                                            buf[16..16 + protected_len].copy_from_slice(&protected[..protected_len]);
                                            message_length = 16 + protected_len;
                                        }
                                        Some(Err(e)) => {
                                            error!("E2E protect error: {:?}", e);
                                        }
                                        None => unreachable!("contains_key was true"),
                                    }
                                }
                            }

                            match socket.send_to(&buf[..message_length], send_message.target_addr).await {
                                Ok(_bytes_sent) => {
                                    trace!("Sent {} bytes to {}", message_length, send_message.target_addr);
                                    if let Ok(()) = send_message.response.send(Ok(())) {} else {
                                        info!("Socket owner closed channel, closing socket.");
                                        // The sender has been dropped, so we should exit
                                        break;
                                    }
                                }
                                Err(e) => {
                                    error!("Failed to send message with error: {:?}", e);
                                    if let Ok(()) = send_message.response.send(Err(Error::Io(e))) {  } else {
                                        error!("Socket owner closed channel unexpectedly, closing socket.");
                                        break;
                                    }
                                }
                            }
                        } else {
                            info!("Send channel closed, closing socket.");
                            // The sender has been dropped, so we should exit
                            break;
                        }
                    }
                }
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::sd::test_support::{TestPayload, empty_sd_header};
    use std::format;

    type TestSocketManager = SocketManager<TestPayload>;

    fn test_registry() -> Arc<Mutex<E2ERegistry>> {
        Arc::new(Mutex::new(E2ERegistry::new()))
    }

    #[tokio::test]
    async fn test_bind_ephemeral_port() {
        let sm = TestSocketManager::bind(0, test_registry()).unwrap();
        assert!(sm.port() > 0);
        assert_eq!(sm.session_id(), 0);
    }

    #[tokio::test]
    async fn test_send_message_new() {
        let target = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 1234);
        let msg = Message::new_sd(1, &empty_sd_header());
        let (rx, send_msg) = SendMessage::<TestPayload>::new(target, msg);
        assert_eq!(send_msg.target_addr, target);
        // Verify the oneshot channel works
        send_msg.response.send(Ok(())).unwrap();
        assert!(rx.await.unwrap().is_ok());
    }

    #[tokio::test]
    async fn test_socket_manager_shut_down() {
        let sm = TestSocketManager::bind(0, test_registry()).unwrap();
        sm.shut_down().await;
    }

    #[tokio::test]
    async fn test_socket_manager_send_and_receive() {
        let mut sm = TestSocketManager::bind(0, test_registry()).unwrap();
        let sm_port = sm.port();

        // Create a raw UDP socket to send data to the SocketManager
        let raw_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();

        // Build and encode an SD message
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        let mut buf = vec![0u8; 128];
        let n = msg.encode(&mut buf.as_mut_slice()).unwrap();

        // Send raw bytes to the SocketManager's port
        raw_socket
            .send_to(&buf[..n], SocketAddrV4::new(Ipv4Addr::LOCALHOST, sm_port))
            .await
            .unwrap();

        // Receive the decoded message from the SocketManager
        let result = tokio::time::timeout(std::time::Duration::from_secs(2), sm.receive())
            .await
            .expect("Timed out waiting for message");

        let received = result.unwrap().unwrap();
        assert_eq!(
            received.message.header().message_id(),
            msg.header().message_id()
        );
        assert!(received.message.is_sd());
    }

    #[tokio::test]
    async fn test_poll_receive() {
        let mut sm = TestSocketManager::bind(0, test_registry()).unwrap();
        let sm_port = sm.port();

        // Send a message to the socket manager from a raw socket
        let raw_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        let mut buf = vec![0u8; 128];
        let n = msg.encode(&mut buf.as_mut_slice()).unwrap();
        raw_socket
            .send_to(&buf[..n], SocketAddrV4::new(Ipv4Addr::LOCALHOST, sm_port))
            .await
            .unwrap();

        // Use poll_fn to exercise poll_receive
        let result = tokio::time::timeout(std::time::Duration::from_secs(2), async {
            std::future::poll_fn(|cx| sm.poll_receive(cx)).await
        })
        .await
        .expect("Timed out waiting for poll_receive");

        let received = result.unwrap().unwrap();
        assert!(received.message.is_sd());
    }

    #[tokio::test]
    async fn test_send_drops_when_socket_loop_exits() {
        let mut sm = TestSocketManager::bind(0, test_registry()).unwrap();
        // Shut down the socket loop by dropping the internal channels
        // We can't directly kill the loop, but we can test the error path
        // by sending to a socket manager that has been shut down.
        let port = sm.port();
        assert!(port > 0);

        // Send a valid message first to verify normal operation
        let raw_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let raw_port = raw_socket.local_addr().unwrap().port();
        let target = SocketAddrV4::new(Ipv4Addr::LOCALHOST, raw_port);
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        sm.send(target, msg).await.unwrap();
        assert_eq!(sm.session_id(), 1);

        // Second send increments session
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        sm.send(target, msg).await.unwrap();
        assert_eq!(sm.session_id(), 2);
    }

    #[tokio::test]
    async fn test_received_message_debug() {
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        let received = ReceivedMessage {
            message: msg,
            source: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 5000),
            e2e_status: None,
        };
        let s = format!("{received:?}");
        assert!(s.contains("ReceivedMessage"));
    }

    #[tokio::test]
    async fn test_send_message_debug() {
        let target = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 1234);
        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        let (_rx, send_msg) = SendMessage::<TestPayload>::new(target, msg);
        let s = format!("{send_msg:?}");
        assert!(s.contains("SendMessage"));
    }

    #[tokio::test]
    async fn test_socket_manager_debug() {
        let sm = TestSocketManager::bind(0, test_registry()).unwrap();
        let s = format!("{sm:?}");
        assert!(s.contains("SocketManager"));
        sm.shut_down().await;
    }

    #[tokio::test]
    async fn test_socket_manager_send_to_target() {
        let mut sm = TestSocketManager::bind(0, test_registry()).unwrap();

        // Create a raw socket to receive
        let raw_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let raw_port = raw_socket.local_addr().unwrap().port();

        let msg = Message::<TestPayload>::new_sd(1, &empty_sd_header());
        let target = SocketAddrV4::new(Ipv4Addr::LOCALHOST, raw_port);

        sm.send(target, msg.clone()).await.unwrap();
        assert_eq!(sm.session_id(), 1);

        // Verify the raw socket received data
        let mut recv_buf = vec![0u8; 1400];
        let (len, _addr) = tokio::time::timeout(
            std::time::Duration::from_secs(2),
            raw_socket.recv_from(&mut recv_buf),
        )
        .await
        .expect("Timed out waiting for sent data")
        .unwrap();

        // Decode and verify
        let view = MessageView::parse(&recv_buf[..len]).unwrap();
        assert_eq!(
            view.header().to_owned().message_id(),
            msg.header().message_id()
        );
    }
}