socket_client.cpp

#include "socket_client.h"
#include <cstring>
#include <cerrno>
#include <iostream>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#endif

using namespace std;

static void close_sock(int& socket) {
    if (socket != -1) {
#ifdef _WIN32
        closesocket(socket);
#else
        close(socket);
#endif
        socket = -1;
    }
}

inline void initialize() {
#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
    }
#endif
}

inline void cleanup() {
#ifdef _WIN32
    WSACleanup();
#endif
}

SocketClient::SocketClient(
    int selectTimeout, int recvTimeout
) : selectTimeout(selectTimeout), recvTimeout(recvTimeout), m_socket(-1) {
    initialize();
}

SocketClient::~SocketClient() {
    disconnect();
    cleanup();
}

#ifdef _WIN32
inline string net_error() {
    return to_string(WSAGetLastError());
}
#else
inline string net_error() {
    return strerror(errno);
}
#endif

bool SocketClient::connectToServer(const std::string& host, uint16_t port) {
    std::unique_lock lock(sock_mtx);

    if (m_socket != -1) return true;

    // 创建socket
    m_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (m_socket == -1) {
        std::cerr << "socket() failed: " << net_error() << std::endl;
        return false;
    }

    // 解析主机名
    struct hostent* server = gethostbyname(host.c_str());
    if (server == nullptr) {
        std::cerr << "gethostbyname() failed: " << net_error() << std::endl;
        close_sock(m_socket);
        return false;
    }

    // 设置服务器地址
    sockaddr_in serv_addr{};
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    memcpy(&serv_addr.sin_addr.s_addr, server->h_addr, server->h_length);
    serv_addr.sin_port = htons(port);

#ifdef _WIN32
    DWORD timeout = recvTimeout;
    if (setsockopt(m_socket, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof(timeout)) == SOCKET_ERROR) {
        std::cerr << "setsockopt(SO_RCVTIMEO) failed: " << net_error() << std::endl;
        close_sock(m_socket);
        return false;
    }
#else
    struct timeval timeout;
    timeout.tv_sec = recvTimeout / 1000;
    timeout.tv_usec = (recvTimeout % 1000) * 1000;
    if (setsockopt(m_socket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
        std::cerr << "setsockopt(SO_RCVTIMEO) failed: " << net_error() << std::endl;
        close_sock(m_socket);
        return false;
    }
#endif

    // 连接服务器
    if (connect(m_socket, (sockaddr*)&serv_addr, sizeof(serv_addr)) < 0) {
        std::cerr << "connect() failed: " << net_error() << std::endl;
        close_sock(m_socket);
        return false;
    }

    return true;
}

bool SocketClient::sendData(const void* data, size_t length) {
    bool ok = true;
    {
        std::shared_lock lock(sock_mtx);
        if (m_socket == -1) return false;
        std::lock_guard lock2(send_mtx);

        const char* ptr = static_cast<const char*>(data);
        size_t totalSent = 0;

        while (totalSent < length) {
            int sent = send(m_socket, ptr + totalSent, length - totalSent, 0);
            if (sent <= 0) {
                std::cerr << "send() failed: " << net_error() << std::endl;
                ok = false;
                break;
            }
            totalSent += sent;
        }
    }

    if (!ok) {
        disconnect();
    }

    return ok;
}

bool SocketClient::isDataAvailable() {
    std::shared_lock lock(sock_mtx);
    if (m_socket == -1) return false;

    fd_set readSet;
    FD_ZERO(&readSet);
    FD_SET(m_socket, &readSet);

    timeval timeout{};
    timeout.tv_sec = selectTimeout / 1000;
    timeout.tv_usec = selectTimeout % 1000 * 1000;

    int result = select(m_socket + 1, &readSet, nullptr, nullptr, &timeout);
    if (result > 0) {
        return FD_ISSET(m_socket, &readSet);
    }

    return false;
}

bool SocketClient::receiveData(void* buffer, size_t length) {
    bool ok = true;
    {
        std::shared_lock lock(sock_mtx);
        if (m_socket == -1) return false;

        char* ptr = static_cast<char*>(buffer);
        size_t totalReceived = 0;

        while (totalReceived < length) {
            int received = recv(m_socket, ptr + totalReceived, length - totalReceived, 0);
            if (received <= 0) {
                std::cerr << "recv() failed: " << net_error() << std::endl;
                ok = false;
                break;
            }
            totalReceived += received;
        }
    }

    if (!ok) {
        disconnect();
    }

    return true;
}

void SocketClient::disconnect() {
    std::unique_lock lock(sock_mtx);
    close_sock(m_socket);
}

bool SocketClient::isConnected() const {
    return m_socket != -1;
}