client2.cpp

#include <arpa/inet.h>
#include <cerrno>
#include <cstdint>
#include <cstring>
#include <fcntl.h>
#include <format>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

namespace globals {
constexpr int error_val = -1;
constexpr int k_max_msg = 32 << 20;
} // namespace globals

[[noreturn]] static void die(const std::string_view msg) {
    throw std::runtime_error(std::format("[{}] {}", errno, msg));
}

static void fd_set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) die("fcntl F_GETFL");
    if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) die("fcntl F_SETFL");
}

class Client {
public:
    Client(const std::string& host, uint16_t port) {
        sock_fd_ = socket(AF_INET, SOCK_STREAM, 0);
        if (sock_fd_ == globals::error_val) die("socket");

        sockaddr_in addr{};
        addr.sin_family = AF_INET;
        addr.sin_port = htons(port);
        if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) <= 0) {
            die("inet_pton");
        }

        fd_set_nonblocking(sock_fd_);

        int ret = connect(sock_fd_, reinterpret_cast<sockaddr*>(&addr), sizeof(addr));
        if (ret == -1 && errno != EINPROGRESS) {
            close(sock_fd_);
            die(std::format("connect failed: {}", strerror(errno)));
        }

        if (ret == -1 && errno == EINPROGRESS) {
            pollfd pfd = {sock_fd_, POLLOUT, 0};
            if (poll(&pfd, 1, 5000) == -1) {
                close(sock_fd_);
                die("poll for connect");
            }
            if (pfd.revents & POLLOUT) {
                int err;
                socklen_t len = sizeof(err);
                if (getsockopt(sock_fd_, SOL_SOCKET, SO_ERROR, &err, &len) == -1) {
                    close(sock_fd_);
                    die("getsockopt");
                }
                if (err != 0) {
                    close(sock_fd_);
                    die(std::format("connection failed: {}", strerror(err)));
                }
            } else {
                close(sock_fd_);
                die("connection timeout");
            }
        }
        std::cerr << "Connected to server\n";
    }

    ~Client() {
        if (sock_fd_ != -1) close(sock_fd_);
    }

    void send_message(const std::string& msg) {
        if (msg.size() > globals::k_max_msg) {
            std::cerr << "Message too long\n";
            return;
        }
        uint32_t len = static_cast<uint32_t>(msg.size());
        outgoing_.insert(outgoing_.end(), reinterpret_cast<uint8_t*>(&len), 
                        reinterpret_cast<uint8_t*>(&len) + 4);
        outgoing_.insert(outgoing_.end(), msg.begin(), msg.end());
        want_write_ = true;
    }

    void handle_events() {
        pollfd pfd = {sock_fd_, 0, 0};
        if (want_read_) pfd.events |= POLLIN;
        if (want_write_) pfd.events |= POLLOUT;

        if (poll(&pfd, 1, -1) == -1) {
            if (errno == EINTR) return;
            die("poll");
        }

        if (pfd.revents & POLLIN) handle_read();
        if (pfd.revents & POLLOUT) handle_write();
        if (pfd.revents & POLLERR) {
            std::cerr << "Socket error\n";
            want_close_ = true;
        }
    }

    bool should_close() const { return want_close_; }

    const std::vector<std::string>& received_messages() const { return messages_; }

private:
    int sock_fd_;
    bool want_read_ = true;
    bool want_write_ = false;
    bool want_close_ = false;
    std::vector<uint8_t> incoming_;
    std::vector<uint8_t> outgoing_;
    std::vector<std::string> messages_;

    void handle_read() {
        std::vector<uint8_t> buffer(64 * 1024);
        ssize_t n_bytes = read(sock_fd_, buffer.data(), buffer.size());
        if (n_bytes == -1) {
            if (errno != EAGAIN) {
                std::cerr << std::format("read() error: {}\n", strerror(errno));
                want_close_ = true;
            }
            return;
        }
        if (n_bytes == 0) {
            std::cerr << "Server closed connection\n";
            want_close_ = true;
            return;
        }

        incoming_.insert(incoming_.end(), buffer.begin(), buffer.begin() + n_bytes);
        while (try_process_response()) {}
    }

    void handle_write() {
        if (outgoing_.empty()) return;
        ssize_t n_bytes = write(sock_fd_, outgoing_.data(), outgoing_.size());
        if (n_bytes == -1) {
            if (errno != EAGAIN) {
                std::cerr << std::format("write() error: {}\n", strerror(errno));
                want_close_ = true;
            }
            return;
        }
        outgoing_.erase(outgoing_.begin(), outgoing_.begin() + n_bytes);
        if (outgoing_.empty()) want_write_ = false;
    }

    bool try_process_response() {
        if (incoming_.size() < 4) return false;

        uint32_t len;
        memcpy(&len, incoming_.data(), 4);
        if (len > globals::k_max_msg) {
            std::cerr << "Response too long\n";
            want_close_ = true;
            return false;
        }

        if (incoming_.size() < 4 + len) return false;

        std::string msg(reinterpret_cast<char*>(incoming_.data() + 4), len);
        messages_.push_back(msg);
        std::cout << std::format("Received: {}\n", msg);

        incoming_.erase(incoming_.begin(), incoming_.begin() + 4 + len);
        return true;
    }
};

int main() {
    try {
        Client client("127.0.0.1", 1234);
        std::string input;

        while (!client.should_close()) {
            std::cout << "Enter message (or 'quit' to exit): ";
            if (std::getline(std::cin, input)) {
                if (input == "quit") break;
                if (!input.empty()) {
                    client.send_message(input);
                }
            }
            client.handle_events();
        }

        for (const auto& msg : client.received_messages()) {
            std::cout << std::format("Echoed: {}\n", msg);
        }
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }
    return 0;
}