Crafter.Review/tests/PgWire/main.cpp

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

// The SQL helper: its statements, refusals, the URL, SCRAM, the protocol on a
// stand-in server (FakePg, tests/Common). Then what the helper prints over
// the same stand-in, line for line against the goldens (fixture/pipeline.json,
// `runs`), and the broker: the URL stays in the driver, a pass gets an
// address and a token.

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
import std;
import Crafter.Review;
namespace fs = std::filesystem;
using namespace Crafter::Review;

#include "../Common/FakePg.hpp"
using namespace Testing;

namespace {
    std::int32_t Failures = 0;

    void Check(std::string_view name, bool cond, std::string_view detail = {}) {
        std::println("{}{}{}", cond ? "  ok   " : "  FAIL ", name, !cond && !detail.empty() ? std::format("  ({})", detail) : std::string());
        if (!cond) ++Failures;
    }


    std::uint16_t DeadPort() {
        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);
        bind(s, reinterpret_cast<sockaddr*>(&address), sizeof(address));
        socklen_t length = sizeof(address);
        getsockname(s, reinterpret_cast<sockaddr*>(&address), &length);
        close(s);
        return ntohs(address.sin_port);
    }

    std::vector<std::string> Of(const std::vector<std::pair<std::string, std::string>>& log, std::string_view kind) {
        std::vector<std::string> out;
        for (const auto& [k, x] : log) {
            if (k == kind) out.push_back(x);
        }
        return out;
    }

    bool Any(const std::vector<std::string>& lines, std::string_view part) {
        return std::ranges::any_of(lines, [&](const std::string& line) { return line.contains(part); });
    }

    std::string Show(const std::vector<std::string>& lines) {
        Json out = Json::array();
        for (const std::string& line : lines) out.push_back(line);
        return Dyn::ReprValue(out);
    }

    // A raw call of the broker: one line in, every line back.
    std::vector<std::string> RawCall(std::uint16_t port, const std::string& line) {
        return Dyn::Split(Dyn::RStrip(BrokerCall(port, line), "\n"), "\n");
    }
}

int main() {
    std::println("the SQL helper (pgwire): statements, refusals, the URL, SCRAM, the protocol on a stand-in server");
    using Pieces = std::vector<std::string>;
    auto sp = [](std::string_view sql) { return PgWire::SplitStatements(sql); };
    Check("split: semicolons in quotes, E-strings, dollar quotes (tagged, and $ inside identifiers), comments (nested, -- to \\r), double quotes and parentheses do not cut; empty statements drop",
          sp("select ';' ; select 2;") == Pieces{"select ';'", "select 2"} && sp("select $x$ a; $y$ ; $x$; commit") == Pieces{"select $x$ a; $y$ ; $x$", "commit"}
          && sp("select 1 as a$b$; commit; select $b$") == Pieces{"select 1 as a$b$", "commit", "select $b$"} && sp("select E'a\\'; commit; '") == Pieces{"select E'a\\'; commit; '"}
          && sp("-- c;\nselect 1; /* x; /* y; */ z; */ select 2") == Pieces{"-- c;\nselect 1", "/* x; /* y; */ z; */ select 2"} && sp("select 1 --x\r; commit") == Pieces{"select 1 --x", "commit"}
          && sp("create rule r as on insert to t do also (insert into a values (1); update b set c = 1); select 3")
                 == Pieces{"create rule r as on insert to t do also (insert into a values (1); update b set c = 1)", "select 3"}
          && sp("select \"a;b\" from t; select 2") == Pieces{"select \"a;b\" from t", "select 2"} && sp("  ;; -- only\n ; ").empty());
    const std::vector<std::string> refuse = {"commit", "COMMIT AND CHAIN", "end", "END WORK", "rollback", "rollback work", "abort", "begin", "start transaction", "prepare transaction 'x'",
                                             "commit prepared 'x'", "rollback prepared 'x'", "/* c */ commit", "-- c\ncommit", "--c\rcommit", "(commit)", "copy t to stdout", "COPY t FROM PROGRAM 'x'"};
    const std::vector<std::string> allow = {"rollback to s1", "ROLLBACK TO SAVEPOINT s", "savepoint s", "release s", "prepare p as select 1", "select 1", "with x as (select 1) select * from x",
                                            "do $$ begin commit; end $$", "call p()", "endless_fn()", "end1()"};
    std::vector<std::string> wrong;
    for (const std::string& s : refuse) {
        if (!PgWire::Refused(s)) wrong.push_back(s);
    }
    for (const std::string& s : allow) {
        if (PgWire::Refused(s)) wrong.push_back(s);
    }
    Check("refused: every statement that ends, restarts or prepares the transaction, and COPY - past comments and parentheses; ROLLBACK TO, SAVEPOINT, DO and CALL go to the server (which refuses COMMIT inside them in a transaction block)",
          wrong.empty(), Show(wrong));
    Check("URL: this machine only; the password masked for a log or a dump", !PgWire::CheckUrl("postgresql://postgres:postgres@127.0.0.1:54322/postgres") && !PgWire::CheckUrl("postgres://u@localhost/x")
          && !PgWire::CheckUrl("postgresql://u@[::1]:5/x") && PgWire::CheckUrl("postgresql://u:p@db.abcdefgh.supabase.co:5432/postgres").value_or("").contains("not this machine")
          && PgWire::CheckUrl("http://127.0.0.1:54321") && PgWire::CheckUrl("") && PgWire::CheckUrl(nullptr)
          && PgWire::Redact("postgresql://postgres:hunter2@127.0.0.1:54322/postgres") == "postgresql://postgres:***@127.0.0.1:54322/postgres"
          && PgWire::Redact("postgresql://127.0.0.1/x") == "postgresql://127.0.0.1/x");
    PgWire::Scram sc("pencil", "user", "rOprNGfwEbeRWgbNEkqO");
    const std::string sf = "r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096";
    Check("SCRAM-SHA-256: the RFC 7677 test vector - the client-first, the client proof, the server signature verified (and a wrong one refused)",
          sc.ClientFirst() == "n,,n=user,r=rOprNGfwEbeRWgbNEkqO"
          && sc.ClientFinal(sf) == "c=biws,r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,p=dHzbZapWIk4jUhN+Ute9ytag9zjfMHgsqmmiz7AndVQ="
          && sc.Verify("v=6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4=") && !sc.Verify("v=AAAATRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4="));
    std::vector<std::string> out;
    const PgWire::Out collect = [&](const std::string& line) { out.push_back(line); };
    {
        FakePg pg;
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "select 1; insert into t values ('a'); boom; select 3", collect);
        const auto log = pg.Served();
        const std::vector<std::string> queries = Of(log, "Q");
        const std::vector<std::string> setup(queries.begin(), queries.begin() + std::min<std::ptrdiff_t>(4, static_cast<std::ptrdiff_t>(queries.size())));
        Check("run: the password, BEGIN and the three SET LOCALs, then one Parse per statement up to the failing one, then ROLLBACK and Terminate; the call says it connected, prints rows, the notice, the error and that it stopped; exit 1",
              !log.empty() && log[0] == std::pair<std::string, std::string>{"password", "pw"}
              && setup == std::vector<std::string>{"BEGIN", "SET LOCAL statement_timeout = '20s'", "SET LOCAL lock_timeout = '3s'", "SET LOCAL idle_in_transaction_session_timeout = '60s'"}
              && Of(log, "P") == std::vector<std::string>{"select 1", "insert into t values ('a')", "boom"} && log.size() >= 2 && log[log.size() - 2] == std::pair<std::string, std::string>{"Q", "ROLLBACK"}
              && log.back() == std::pair<std::string, std::string>{"X", ""} && code == 1 && out.at(0).starts_with(PgWire::Answered) && std::ranges::contains(out, std::string("NOTICE: hello"))
              && Any(out, "syntax error") && Any(out, "statement 3 failed") && out.back() == "-- rolled back: nothing of this call was committed", Show(out));
    }
    {
        FakePg pg;
        std::unique_ptr<PgWire::Connection> cn = PgWire::Connect(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), 1);
        cn->Simple("BEGIN");
        const PgWire::Answer two = cn->Extended("select 1; select 2");
        cn->Simple("ROLLBACK");
        cn->Close();
        const auto log = pg.Served();
        Check("extended: a whole script in one Parse comes back as the server refuses it (two commands), the transaction failed - never run", two.error && two.error->at('C') == "42601"
              && two.rows.empty() && std::ranges::contains(log, std::pair<std::string, std::string>{"P", "select 1; select 2"}));
    }
    {
        out.clear();
        FakePg pg;
        // A statement the refusal list misses: the transaction status is the
        // backstop.
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "create table x (a int); commit; select 1", collect, {},
                                                 [](std::string_view) { return std::optional<std::string>(); });
        const auto log = pg.Served();
        Check("run: the backstop - after a statement that ended the transaction nothing more is sent, the call fails and says what may have been committed",
              Of(log, "P") == std::vector<std::string>{"create table x (a int)", "commit"} && code == 1 && Any(out, "STOPPED") && out.back().contains("may have been committed"), Show(out));
    }
    {
        out.clear();
        FakePg pg;
        const std::string url = std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port);
        const std::int32_t code = PgWire::RunSql(url, "select 1; commit", collect);
        PgWire::Connect(url, 1)->Close();
        const auto log = pg.Served();
        Check("run: a refused statement stops the call before it connects - nothing is sent, no connected line", code == 1 && out.at(0).starts_with("review-gates sql: refused, nothing ran")
              && out.at(0).contains("statement 2: COMMIT") && !Any(out, PgWire::Answered) && Of(log, "Q").empty() && Of(log, "P").empty());
    }
    {
        out.clear();
        const std::uint16_t dead = DeadPort();
        const PgWire::Connector once = [](std::string_view url) { return PgWire::Connect(url, 1, [](double) {}); };
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", dead), "select 1", collect, once);
        std::vector<std::string> ignored;
        const PgWire::Out drop = [&](const std::string& line) { ignored.push_back(line); };
        Check("run: nothing listening -> NO DATABASE, exit 2, no connected line; a hosted URL never connects; no statement is a usage error", code == 2 && out.at(0).starts_with(PgWire::NoDatabase)
              && PgWire::RunSql("postgresql://u@db.x.supabase.co/x", "select 1", drop) == 2 && PgWire::RunSql(std::format("postgresql://u@127.0.0.1:{}/x", dead), " -- nothing\n", drop) == 64);
        Check("...and the reason: cannot connect to 127.0.0.1:<port> ([Errno 111] Connection refused)",
              out.at(0) == std::format("{} - cannot connect to 127.0.0.1:{} ([Errno 111] Connection refused)", PgWire::NoDatabase, dead), out.at(0));
    }
    {
        out.clear();
        FakePg pg("other");
        const PgWire::Connector once = [](std::string_view url) { return PgWire::Connect(url, 1); };
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), "select 1", collect, once);
        const std::string detail = out.at(0).substr(out.at(0).find(" - ") + 3);
        Check("run: a refused login is NO DATABASE (exit 2), its message without the password", code == 2 && out.at(0).starts_with(PgWire::NoDatabase) && out.at(0).contains("28P01")
              && !detail.contains("pw"), Show(out));
    }
    {
        // A busy server is asked again, but not slept on after the last try.
        std::vector<double> slept;
        try {
            PgWire::Connect(std::format("postgresql://u@127.0.0.1:{}/x", DeadPort()), 3, [&](double s) { slept.push_back(s); });
        } catch (const PgWire::DbError&) {
            // Expected: nothing listens.
        }
        Check("connect: an unreachable server is tried three times, slept on between tries only", slept == std::vector<double>{0.5, 1.0});
    }

    // What the helper prints over the stand-in, line for line against the
    // goldens.
    const Json goldens = Dyn::LoadJsonFile(fs::current_path() / "tests" / "Goldens" / "fixture" / "pipeline.json");
    std::int32_t compared = 0;
    for (const Json& entry : goldens.at("pgwire").at("runs")) {
        if (entry.contains("noServer")) continue;
        out.clear();
        FakePg pg;
        const std::int32_t code = PgWire::RunSql(std::format("postgresql://u:pw@127.0.0.1:{}/db", pg.port), entry.at("sql").get<std::string>(), collect);
        pg.Served();
        std::vector<std::string> normalized;
        for (const std::string& line : out) normalized.push_back(Dyn::Replace(line, std::format(":{}", pg.port), ":PORT"));
        std::vector<std::string> want;
        for (const Json& line : entry.at("out")) want.push_back(line.get<std::string>());
        if (code == entry.at("code").get<std::int32_t>() && normalized == want) ++compared;
        else Check(std::format("run over the stand-in, as the goldens record it: {}", entry.at("sql").get<std::string>()), false, Show(normalized));
    }
    Check(std::format("run over the stand-in prints what the goldens record, line for line ({} scripts)", compared), compared == 6);

    // The broker.
    std::vector<std::string> ran;
    PgWire::Broker broker(
        [&](const std::string& url, const std::string& sql, const PgWire::Out& say) {
            ran.push_back(std::format("{}|{}", url, sql));
            say(std::format("{} to x", PgWire::Answered));
            say("é line");
            return 1;
        },
        [](const std::string& url) { return std::pair<bool, std::string>{url.contains("good"), "detail"}; });
    broker.Start();
    const std::string url = "postgresql://postgres:hunter2@127.0.0.1:54322/postgres";
    const std::string address = broker.Address(url);
    const std::string token = address.substr(address.rfind('#') + 1);
    Check("broker: the address is 127.0.0.1:<port>#<32 hex>, one token per URL, and the token names its URL", address.starts_with("127.0.0.1:") && token.size() == 32
          && std::ranges::all_of(token, [](char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'); }) && broker.Address(url) == address
          && broker.Address("postgresql://other@127.0.0.1/x") != address && broker.UrlOf(token) == url && broker.Address("").empty() && !address.contains("hunter2"));
    out.clear();
    std::int32_t code = PgWire::CallBroker(address, "select 1;\nselect 'é';", false, collect);
    Check("broker: a call runs the SQL on its token's URL, streams the lines back and returns the call's exit code", code == 1 && ran.size() == 1 && ran[0] == std::format("{}|select 1;\nselect 'é';", url)
          && out == std::vector<std::string>{std::format("{} to x", PgWire::Answered), "é line"}, Show(out));
    out.clear();
    code = PgWire::CallBroker(broker.Address("postgresql://good@127.0.0.1/x"), "", true, collect);
    Check("broker: a probe answers with the connected line and exit 0", code == 0 && out == std::vector<std::string>{std::format("{} - detail", PgWire::Answered)}, Show(out));
    out.clear();
    code = PgWire::CallBroker(std::format("{}#{}", address.substr(0, address.rfind('#')), std::string(32, '0')), "select 1", false, collect);
    Check("broker: an unknown token is NO DATABASE, exit 2, and nothing runs", code == 2 && out.size() == 1 && out[0].starts_with(PgWire::NoDatabase) && ran.size() == 1, Show(out));
    const std::uint16_t brokerPort = static_cast<std::uint16_t>(std::stoi(address.substr(10, address.rfind('#') - 10)));
    const std::vector<std::string> notWhole = RawCall(brokerPort, "{\"token\": \"x\"");
    Check("broker: a call that never ends its line did not arrive whole - nothing runs, exit 64", notWhole.size() == 2 && notWhole[0].contains("did not arrive whole")
          && notWhole[1] == std::format("{}64", PgWire::ExitMark), Show(notWhole));
    const std::vector<std::string> listToken = RawCall(brokerPort, "{\"token\": [1], \"sql\": \"select 1\"}\n");
    Check("broker: a token that is no string gets NO DATABASE and the exit mark", listToken.size() == 2
          && listToken[0].starts_with(PgWire::NoDatabase) && listToken[1] == std::format("{}2", PgWire::ExitMark), Show(listToken));
    out.clear();
    Check("call_broker: no address, a bad port, a closed port - NO DATABASE, exit 2, with the reason",
          PgWire::CallBroker("", "x", false, collect) == 2 && PgWire::CallBroker("127.0.0.1:x#t", "x", false, collect) == 2
          && PgWire::CallBroker(std::format("127.0.0.1:{}#t", DeadPort()), "x", false, collect) == 2 && out.size() == 3
          && out[0] == std::format("{} - this pass has no database (the driver gives one to a pass with the supabase environment)", PgWire::NoDatabase)
          && out[1] == std::format("{} - the driver's SQL broker does not answer (invalid literal for int() with base 10: 'x')", PgWire::NoDatabase)
          && out[2].ends_with("does not answer ([Errno 111] Connection refused)"), Show(out));
    broker.Close();
    out.clear();
    Check("broker: closed, its port answers no more", PgWire::CallBroker(address, "select 1", false, collect) == 2 && out.at(0).ends_with("([Errno 111] Connection refused)"), Show(out));
    return Failures ? 1 : 0;
}