Crafter.Review/tests/Common/FakePg.hpp

// SPDX-License-Identifier: AGPL-3.0-only
// SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// A Postgres stand-in on 127.0.0.1 for the tests that talk the wire protocol
// (PgWire, ReviewFixes), and a raw call of the broker. Included after
// `import std;`, `import Crafter.Review;` and the socket headers.
namespace Testing {
    std::string Be32(std::uint32_t v) {
        return {static_cast<char>(v >> 24), static_cast<char>((v >> 16) & 0xFF), static_cast<char>((v >> 8) & 0xFF), static_cast<char>(v & 0xFF)};
    }

    std::string Msg(char type, const std::string& body = {}) {
        return std::format("{}{}{}", type, Be32(static_cast<std::uint32_t>(body.size() + 4)), body);
    }

    std::string Err(const std::string& code, const std::string& text) {
        return Msg('E', std::format("SERROR{}C{}{}M{}{}{}", std::string(1, '\0'), code, std::string(1, '\0'), text, std::string(1, '\0'), std::string(1, '\0')));
    }

    // A raw call of the broker on 127.0.0.1: the bytes as they are, the
    // write side shut, everything that comes back.
    std::string BrokerCall(std::uint16_t port, std::string_view bytes) {
        const std::int32_t s = socket(AF_INET, SOCK_STREAM, 0);
        sockaddr_in address {};
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        address.sin_port = htons(port);
        connect(s, reinterpret_cast<sockaddr*>(&address), sizeof(address));
        ::send(s, bytes.data(), bytes.size(), MSG_NOSIGNAL);
        shutdown(s, SHUT_WR);
        std::string all;
        std::array<char, 4096> chunk {};
        ssize_t got = 0;
        while ((got = recv(s, chunk.data(), chunk.size(), 0)) > 0) all.append(chunk.data(), static_cast<std::size_t>(got));
        close(s);
        return all;
    }

    // A Postgres stand-in on 127.0.0.1 for one connection: cleartext
    // password, simple Query, Parse/Bind/Describe/Execute/Sync. It records
    // what it is sent and answers as the server does where the helper relies
    // on it: two commands in one Parse are an error; COMMIT ends the
    // transaction (status I); a statement after an error fails.
    class FakePg {
    public:
        explicit FakePg(std::string password = "pw") : password(std::move(password)) {
            listener = socket(AF_INET, SOCK_STREAM, 0);
            sockaddr_in address {};
            address.sin_family = AF_INET;
            address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
            bind(listener, reinterpret_cast<sockaddr*>(&address), sizeof(address));
            listen(listener, 1);
            socklen_t length = sizeof(address);
            getsockname(listener, reinterpret_cast<sockaddr*>(&address), &length);
            port = ntohs(address.sin_port);
            thread = std::thread([this] { Serve(); });
        }

        ~FakePg() {
            Served();
        }

        // The log once the connection is over.
        std::vector<std::pair<std::string, std::string>> Served() {
            if (thread.joinable()) thread.join();
            return log;
        }

        std::uint16_t port = 0;

    private:
        std::string password;
        std::int32_t listener = -1;
        std::thread thread;
        std::vector<std::pair<std::string, std::string>> log;

        void Serve() {
            const std::int32_t c = accept(listener, nullptr, nullptr);
            std::string buffer;
            auto read = [&](std::size_t n) {
                while (buffer.size() < n) {
                    std::array<char, 4096> chunk {};
                    const ssize_t got = recv(c, chunk.data(), chunk.size(), 0);
                    if (got <= 0) break;
                    buffer.append(chunk.data(), static_cast<std::size_t>(got));
                }
                const std::string out = buffer.substr(0, std::min(n, buffer.size()));
                buffer.erase(0, out.size());
                return out;
            };
            auto length = [](const std::string& bytes) {
                return (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[0])) << 24) | (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[1])) << 16)
                       | (static_cast<std::uint32_t>(static_cast<std::uint8_t>(bytes[2])) << 8) | static_cast<std::uint8_t>(bytes[3]);
            };
            auto send = [&](const std::string& data) { ::send(c, data.data(), data.size(), MSG_NOSIGNAL); };
            const std::uint32_t startup = length(read(4));
            read(startup - 4);
            send(Msg('R', Be32(3)));
            read(1);
            std::string got = read(length(read(4)) - 4);
            while (!got.empty() && got.back() == '\0') got.pop_back();
            log.emplace_back("password", got);
            if (got != password) {
                send(Err("28P01", "password authentication failed"));
                close(c);
                close(listener);
                return;
            }
            send(Msg('S', std::string("server_version\x00" "17.6\x00", 20)) + Msg('R', Be32(0)) + Msg('Z', "I"));
            std::string status = "I";
            std::string pending;
            while (true) {
                const std::string type = read(1);
                if (type.empty()) break;
                const std::string body = read(length(read(4)) - 4);
                if (type == "Q") {
                    std::string q = body;
                    while (!q.empty() && q.back() == '\0') q.pop_back();
                    log.emplace_back("Q", q);
                    status = q.starts_with("BEGIN") ? "T" : (q == "ROLLBACK" ? "I" : status);
                    send(Msg('C', Dyn::Split(q)[0] + std::string(1, '\0')) + Msg('Z', status));
                } else if (type == "P") {
                    pending = body.substr(1, body.find('\0', 1) - 1);
                    log.emplace_back("P", pending);
                } else if (type == "S") {
                    const std::string low = Dyn::Lower(pending);
                    std::string out;
                    if (status == "E") {
                        out = Err("25P02", "current transaction is aborted");
                    } else if (pending.contains(';')) {
                        out = Err("42601", "cannot insert multiple commands into a prepared statement");
                        status = "E";
                    } else if (low.starts_with("commit")) {
                        out = Msg('C', std::string("COMMIT\0", 7));
                        status = "I";
                    } else if (low.starts_with("select")) {
                        out = Msg('T', std::string("\x00\x01x\x00", 4) + std::string(18, '\0')) + Msg('D', std::string("\x00\x01\x00\x00\x00\x01" "1", 7))
                              + Msg('N', std::string("SNOTICE\0Mhello\0\0", 16)) + Msg('C', std::string("SELECT 1\0", 9));
                    } else if (low.starts_with("boom")) {
                        out = Err("42601", "syntax error at or near \"boom\"");
                        status = "E";
                    } else {
                        out = Msg('C', Dyn::Upper(Dyn::Split(pending)[0]) + std::string(1, '\0'));
                    }
                    send(Msg('1') + Msg('2') + out + Msg('Z', status));
                } else if (type == "X") {
                    log.emplace_back("X", "");
                    break;
                }
            }
            close(c);
            close(listener);
        }
    };
}