networks-2.rs/network/windows/src/eventloop/socket.rs

use crate::address::WindowsAddressKnown;
use crate::errors::throw_from_windows_err_code;
use crate::eventloop::allocator::Allocator;
use crate::eventloop::task::IoTask;
use std::mem::uninitialized;
use std::pin::Pin;
use std::ptr::NonNull;
use std::slice;
use windows::Win32::Networking::WinSock::{
    AF_UNSPEC, IPPROTO_ICMP, LPWSAOVERLAPPED_COMPLETION_ROUTINE, SOCK_RAW, SOCKADDR,
    SOCKADDR_STORAGE, SOCKET, SOCKET_ERROR, WSA_FLAG_NO_HANDLE_INHERIT, WSA_FLAG_OVERLAPPED,
    WSA_IO_PENDING, WSA_OPERATION_ABORTED, WSAEACCES, WSAEADDRNOTAVAIL, WSAEAFNOSUPPORT,
    WSAECONNRESET, WSAEDESTADDRREQ, WSAEFAULT, WSAEHOSTUNREACH, WSAEINPROGRESS, WSAEINTR,
    WSAEINVAL, WSAEMSGSIZE, WSAENETDOWN, WSAENETRESET, WSAENETUNREACH, WSAENOBUFS, WSAENOTCONN,
    WSAENOTSOCK, WSAESHUTDOWN, WSAEWOULDBLOCK, WSAGetLastError, WSANOTINITIALISED, WSARecvFrom,
    WSASendTo, WSASocketW, closesocket,
};
use windows::Win32::System::IO::OVERLAPPED;

pub(super) struct WindowsOverlappingIcmpSocket<'a> {
    socket: SOCKET,
    heap: &'a Allocator,
}

impl<'a> WindowsOverlappingIcmpSocket<'a> {
    pub fn new(heap: &'a Allocator) -> Self {
        let result = unsafe {
            WSASocketW(
                AF_UNSPEC.0 as i32,
                SOCK_RAW.0,
                IPPROTO_ICMP.0,
                None,
                0,
                WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT,
            )
        };

        match result {
            Ok(s) => return Self { socket: s, heap },
            Err(e) => throw_from_windows_err_code(e.code()),
        }
    }
    unsafe extern "system" fn _native_callback<
        F: for<'xx> FnOnce(u32, &'xx [u8], WindowsAddressKnown) + 'static,
    >(
        dwerror: u32,
        cbtransferred: u32,
        lpoverlapped: *mut OVERLAPPED,
        dwflags: u32,
    ) {
        let task =
            NonNull::from_mut(IoTask::<F>::from_overlapped(lpoverlapped).get_unchecked_mut())
                .as_ptr();

        (task.read().callback)(
            dwerror,
            slice::from_raw_parts(
                NonNull::from_ref(&(*task).buff).as_ptr().cast(),
                cbtransferred as usize,
            ),
            task.read().addr,
        );

        NonNull::new_unchecked(task)
            .as_mut()
            .heap
            .as_mut()
            .free(Pin::new_unchecked(NonNull::new_unchecked(task).as_mut()));
    }

    fn _gen_native_callback<F: for<'xx> FnOnce(u32, &'xx [u8], WindowsAddressKnown) + 'static>(
        _: &IoTask<F>,
    ) -> LPWSAOVERLAPPED_COMPLETION_ROUTINE {
        return Some(Self::_native_callback::<F>);
    }

    pub unsafe fn send(
        &mut self,
        addr: WindowsAddressKnown,
        init_data: impl FnOnce(&mut [u8]) -> usize,
        on_sent: impl FnOnce(u32, WindowsAddressKnown) + 'static,
    ) {
        let task = self.heap.alloc(|p| {
            IoTask::init(
                p,
                NonNull::from_ref(self.heap),
                addr,
                |errcode, _: &[u8], addr| on_sent(errcode, addr),
            )
        });
        let task_p = task.get_unchecked_mut();

        task_p.buff_d[0].len = init_data(&mut task_p.buff) as u32;

        let result = WSASendTo(
            self.socket,
            &task_p.buff_d,
            None,
            0,
            Some(NonNull::from_ref(&task_p.addr).cast().as_ptr()),
            size_of::<SOCKADDR_STORAGE>() as i32,
            Some(NonNull::from_mut(&mut task_p.sys).as_ptr()),
            Self::_gen_native_callback(task_p),
        );

        if result != SOCKET_ERROR {
            return;
        }

        let result = WSAGetLastError();
        match result {
            WSA_IO_PENDING => {}
            WSAENETUNREACH | WSAENETRESET | WSAECONNRESET => todo!("report"),
            WSAEWOULDBLOCK => throw_from_windows_err_code(result),
            WSAEACCES
            | WSAEADDRNOTAVAIL
            | WSAEAFNOSUPPORT
            | WSAEDESTADDRREQ
            | WSAEFAULT
            | WSAEHOSTUNREACH
            | WSAEINPROGRESS
            | WSAEINTR
            | WSAEINVAL
            | WSAEMSGSIZE
            | WSAENETDOWN
            | WSAENOBUFS
            | WSAENOTCONN
            | WSAENOTSOCK
            | WSAESHUTDOWN
            | WSANOTINITIALISED
            | WSA_OPERATION_ABORTED => throw_from_windows_err_code(result),
            _ => throw_from_windows_err_code(result),
        }
    }

    pub unsafe fn receive(
        &self,
        on_received: impl for<'xx> FnOnce(u32, WindowsAddressKnown, &'xx [u8]) + 'static,
    ) {
        let task = self.heap.alloc(|p| {
            IoTask::init(
                p,
                NonNull::from_ref(self.heap),
                uninitialized(),
                |errcode, buffer: &[u8], addr| on_received(errcode, addr, buffer),
            )
        });
        let task_p = task.get_unchecked_mut();

        task_p._addr_size = size_of::<SOCKADDR_STORAGE>() as i32;
        task_p._flags = 0;
        let result = WSARecvFrom(
            self.socket,
            &task_p.buff_d,
            None,
            &mut task_p._flags,
            Some(
                NonNull::from_mut(&mut task_p.addr.any.native)
                    .cast::<SOCKADDR>()
                    .as_ptr(),
            ),
            Some(&mut task_p._addr_size),
            Some(&mut task_p.sys),
            Self::_gen_native_callback(task_p),
        );

        if result != SOCKET_ERROR {
            return;
        }

        let result = WSAGetLastError();
        match result {
            WSAECONNRESET | WSAEMSGSIZE | WSAENETRESET | WSA_IO_PENDING => return,
            WSAEINVAL => throw_from_windows_err_code(result),
            WSAEWOULDBLOCK => throw_from_windows_err_code(result),
            WSAEFAULT
            | WSAEINPROGRESS
            | WSAEINTR
            | WSAENETDOWN
            | WSAENOTCONN
            | WSANOTINITIALISED
            | WSA_OPERATION_ABORTED => throw_from_windows_err_code(result),
            _ => throw_from_windows_err_code(result),
        }
    }
}

impl Drop for WindowsOverlappingIcmpSocket<'_> {
    fn drop(&mut self) {
        unsafe {
            if closesocket(self.socket) == SOCKET_ERROR {
                throw_from_windows_err_code(WSAGetLastError())
            }
        }
    }
}