Socket.cpp

Go to the documentation of this file.

#include "pch.h"
#include "System/String.h"
#include "System/Net/Sockets/Socket.h"
#include "System/Net/Sockets/SocketException.h"
#include "System/ArgumentException.h"
#include <mutex>
#include <cstring>
#include <cstdlib>
 
#if defined(_WIN32)
    #define WIN32_LEAN_AND_MEAN
    #include <winsock2.h>
    #include <ws2tcpip.h>
    #pragma comment(lib, "Ws2_32.lib")
    using SockLen = int;
    #ifndef SD_RECEIVE
        #define SD_RECEIVE 0
        #define SD_SEND 1
        #define SD_BOTH 2
    #endif
#else
    #include <sys/socket.h>
    #include <sys/types.h>
    #include <netinet/in.h>
    #include <arpa/inet.h>
    #include <unistd.h>
    #include <fcntl.h>
    #include <errno.h>
    #include <sys/select.h>
    #define SOCKET int
    #define INVALID_SOCKET -1
    #define SOCKET_ERROR -1
    #define SD_RECEIVE 0
    #define SD_SEND 1
    #define SD_BOTH 2
    using SockLen = socklen_t;
#endif
 
namespace DotNetDupe {
    namespace System {
        namespace Net {
            namespace Sockets {
 
                class SocketImpl {
                public:
                    SOCKET hSocket;
 
                    SocketImpl() : hSocket(INVALID_SOCKET) {}
                    ~SocketImpl() {
                        Cleanup();
                    }
 
                    void Cleanup() {
                        if (hSocket != INVALID_SOCKET) {
#if defined(_WIN32)
                            closesocket(hSocket);
#else
                            close(hSocket);
#endif
                            hSocket = INVALID_SOCKET;
                        }
                    }
                };
 
                static int GetLastErrorCode() {
#if defined(_WIN32)
                    return WSAGetLastError();
#else
                    return errno;
#endif
                }
 
                static void InitializeSockets() {
#if defined(_WIN32)
                    static bool bInitialized = false;
                    static std::mutex mutexInit;
                    std::lock_guard<std::mutex> lock(mutexInit);
                    if (!bInitialized) {
                        WSADATA wsaData;
                        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
                            throw SocketException(GetLastErrorCode(), "WSAStartup failed");
                        }
                        bInitialized = true;
                        std::atexit([]() { WSACleanup(); });
                    }
#endif
                }
 
                static SOCKET CreateNativeSocket(AddressFamily addressFamily, SocketType socketType, ProtocolType protocolType) {
                    int af = (addressFamily == AddressFamily::InterNetworkV6) ? AF_INET6 : AF_INET;
                    int type = (socketType == SocketType::Dgram) ? SOCK_DGRAM : SOCK_STREAM;
                    int proto = (protocolType == ProtocolType::Udp) ? IPPROTO_UDP : IPPROTO_TCP;
#if defined(_WIN32)
                    return socket(af, type, proto);
#else
                    return ::socket(af, type, proto);
#endif
                }
 
                static void InitSockAddrIn(const String& ip, int port, sockaddr_in& addr) {
                    std::memset(&addr, 0, sizeof(addr));
                    addr.sin_family = AF_INET;
                    addr.sin_port = htons(static_cast<u_short>(port));
                    if (ip.IsEmpty() || ip == "0.0.0.0") {
                        addr.sin_addr.s_addr = INADDR_ANY;
                    } else {
                        if (inet_pton(AF_INET, ip.GetRawString(), &addr.sin_addr) != 1) {
                            throw ArgumentException("Invalid IP address.");
                        }
                    }
                }
 
                Socket::Socket(AddressFamily addressFamily, SocketType socketType, ProtocolType protocolType)
                    : m_pImpl(new SocketImpl()) {
                    InitializeSockets();

                    m_pImpl->hSocket = CreateNativeSocket(addressFamily, socketType, protocolType);
                    if (m_pImpl->hSocket == INVALID_SOCKET) {
                        throw SocketException(GetLastErrorCode(), "Failed to create native socket.");
                    }
                }
 
                Socket::Socket(void* pNativeHandle)
                    : m_pImpl(new SocketImpl()) {
                    InitializeSockets();

#if defined(_WIN32)
                    m_pImpl->hSocket = reinterpret_cast<SOCKET>(pNativeHandle);
#else
                    m_pImpl->hSocket = static_cast<int>(reinterpret_cast<intptr_t>(pNativeHandle));
#endif
                }
 
                Socket::~Socket() {}
 
                Socket::Socket(Socket&& other) noexcept : m_pImpl(std::move(other.m_pImpl)) {}
 
                Socket& Socket::operator=(Socket&& other) noexcept {
                    if (this != &other) m_pImpl = std::move(other.m_pImpl);
                    return *this;
                }
 
                void Socket::Bind(const String& ip, int port) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, String("Socket is closed."));

                    sockaddr_in addr;
                    InitSockAddrIn(ip, port, addr);
                    int optval = 1;
#if defined(_WIN32)
                    setsockopt(m_pImpl->hSocket, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char*>(&optval), sizeof(optval));
#else
                    ::setsockopt(m_pImpl->hSocket, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
#endif

                    if (bind(m_pImpl->hSocket, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) == SOCKET_ERROR) {
                        throw SocketException(GetLastErrorCode(), "Failed to bind socket.");
                    }
                }
 
                void Socket::Listen(int backlog) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) {
                        throw SocketException(-1, "Socket is closed.");
                    }

                    if (listen(m_pImpl->hSocket, backlog) == SOCKET_ERROR) {
                        throw SocketException(GetLastErrorCode(), "Failed to listen on socket.");
                    }
                }
 
                SmartPointer<Socket> Socket::Accept() {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) {
                        throw SocketException(-1, "Socket is closed.");
                    }

                    sockaddr_in clientAddr;
                    SockLen clientSize = sizeof(clientAddr);
#if defined(_WIN32)
                    SOCKET clientSocket = accept(m_pImpl->hSocket, reinterpret_cast<sockaddr*>(&clientAddr), &clientSize);
#else
                    int clientSocket = ::accept(m_pImpl->hSocket, reinterpret_cast<sockaddr*>(&clientAddr), &clientSize);
#endif
                    if (clientSocket == INVALID_SOCKET) {
                        throw SocketException(GetLastErrorCode(), "Failed to accept connection.");
                    }
 
                    return SmartPointer<Socket>(new Socket(reinterpret_cast<void*>(clientSocket)), true);
                }
 
                void Socket::Connect(const String& ip, int port) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, "Socket is closed.");

                    sockaddr_in addr;
                    InitSockAddrIn(ip, port, addr);
                    if (connect(m_pImpl->hSocket, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) == SOCKET_ERROR) {
                        throw SocketException(GetLastErrorCode(), "Failed to connect to host.");
                    }
                }
 
                int Socket::Send(const char* buffer, int offset, int size) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, "Socket is closed.");

                    int bytesSent = send(m_pImpl->hSocket, buffer + offset, size, 0);
                    if (bytesSent == SOCKET_ERROR) throw SocketException(GetLastErrorCode(), "Failed to send data.");
                    return bytesSent;
                }
 
                int Socket::Receive(char* buffer, int offset, int size) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, "Socket is closed.");

                    int bytesReceived = recv(m_pImpl->hSocket, buffer + offset, size, 0);
                    if (bytesReceived == SOCKET_ERROR) throw SocketException(GetLastErrorCode(), "Failed to receive data.");
                    return bytesReceived;
                }
 
                int Socket::SendTo(const char* buffer, int offset, int size, const String& ip, int port) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, "Socket is closed.");

                    sockaddr_in addr;
                    InitSockAddrIn(ip, port, addr);
                    int bytesSent = sendto(m_pImpl->hSocket, buffer + offset, size, 0, reinterpret_cast<sockaddr*>(&addr), sizeof(addr));
                    if (bytesSent == SOCKET_ERROR) throw SocketException(GetLastErrorCode(), "Failed to send data to host.");
                    return bytesSent;
                }
 
                static void FormatSockAddrIp(const sockaddr_in& addr, String& ip, int& port) {
                    char ipBuf[INET_ADDRSTRLEN];
                    if (inet_ntop(AF_INET, &addr.sin_addr, ipBuf, sizeof(ipBuf)) != nullptr) {
                        ip = String(ipBuf);
                    }
                    port = ntohs(addr.sin_port);
                }
 
                int Socket::ReceiveFrom(char* buffer, int offset, int size, String& ip, int& port) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) throw SocketException(-1, "Socket is closed.");

                    sockaddr_in addr;
                    SockLen addrSize = sizeof(addr);
                    std::memset(&addr, 0, sizeof(addr));
                    int bytesReceived = recvfrom(m_pImpl->hSocket, buffer + offset, size, 0, reinterpret_cast<sockaddr*>(&addr), &addrSize);
                    if (bytesReceived == SOCKET_ERROR) throw SocketException(GetLastErrorCode(), "Failed to receive data from host.");
 
                    FormatSockAddrIp(addr, ip, port);
                    return bytesReceived;
                }
 
                void Socket::Close() {
                    if (m_pImpl) m_pImpl->Cleanup();
                }
 
                void Socket::Shutdown(SocketShutdown how) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) return;

                    int nativeHow = (how == SocketShutdown::Receive) ? SD_RECEIVE : ((how == SocketShutdown::Send) ? SD_SEND : SD_BOTH);
                    shutdown(m_pImpl->hSocket, nativeHow);
                }
 
                static int SelectSingleSocket(SOCKET hSocket, int microSeconds, SelectMode mode) {
                    fd_set fds;
                    FD_ZERO(&fds);
                    FD_SET(hSocket, &fds);
                    timeval timeout{ microSeconds / 1000000, microSeconds % 1000000 };
                    fd_set* pRead = (mode == SelectMode::SelectRead) ? &fds : nullptr;
                    fd_set* pWrite = (mode == SelectMode::SelectWrite) ? &fds : nullptr;
                    fd_set* pErr = (mode == SelectMode::SelectError) ? &fds : nullptr;
#if defined(_WIN32)
                    return select(0, pRead, pWrite, pErr, &timeout);
#else
                    return select(hSocket + 1, pRead, pWrite, pErr, &timeout);
#endif
                }
 
                bool Socket::Poll(int microSeconds, SelectMode mode) {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET) return false;
                    return SelectSingleSocket(m_pImpl->hSocket, microSeconds, mode) > 0;
                }
 
                static bool IsPeerConnected(SOCKET hSocket) {
                    sockaddr_in peerAddr;
                    SockLen peerLen = sizeof(peerAddr);
                    return (getpeername(hSocket, reinterpret_cast<sockaddr*>(&peerAddr), &peerLen) != SOCKET_ERROR);
                }
 
                static bool IsSocketActivePeek(SOCKET hSocket) {
                    char buf;
#if defined(_WIN32)
                    int result = recv(hSocket, &buf, 1, MSG_PEEK);
                    if (result == 0) return false;
                    return (result != SOCKET_ERROR || WSAGetLastError() == WSAEWOULDBLOCK);
#else
                    int result = ::recv(hSocket, &buf, 1, MSG_PEEK | MSG_DONTWAIT);
                    if (result == 0) return false;
                    return (result != SOCKET_ERROR || (errno == EAGAIN || errno == EWOULDBLOCK));
#endif
                }
 
                bool Socket::Connected() const {
                    if (!m_pImpl || m_pImpl->hSocket == INVALID_SOCKET || !IsPeerConnected(m_pImpl->hSocket)) return false;
                    int sel = SelectSingleSocket(m_pImpl->hSocket, 0, SelectMode::SelectRead);
                    if (sel == SOCKET_ERROR) return false;
                    if (sel == 0) return true;
                    return IsSocketActivePeek(m_pImpl->hSocket);
                }
 
                void* Socket::GetNativeHandle() const {
                    return m_pImpl ? reinterpret_cast<void*>(m_pImpl->hSocket) : nullptr;
                }
 
            }
        }
    }
}