233 lines
13 KiB
JavaScript
233 lines
13 KiB
JavaScript
import assert from "node:assert/strict";
|
|
import test from "node:test";
|
|
import { PostgresConnectionRegistry, PINNED_HEADERS, internalScopeOf } from "../src/connection-registry.js";
|
|
import { rolePassword } from "../src/execution-database.js";
|
|
|
|
const SECRET = "0123456789abcdef0123456789abcdef";
|
|
const EXECUTION = "00000000-0000-0000-0000-000000000001";
|
|
const DATABASE = "exec_00000000_0000_0000_0000_000000000001";
|
|
|
|
const pinned = (overrides = {}) => ({
|
|
"x-postgres-scope": EXECUTION,
|
|
"x-postgres-host": "db.internal",
|
|
"x-postgres-database": "postgres",
|
|
"x-postgres-user": "ops",
|
|
"x-postgres-password": "ops-password",
|
|
"x-postgres-role-secret": SECRET,
|
|
...overrides
|
|
});
|
|
const internal = (id = "mcp-a", overrides) => ({ sessionId: id, pinned: pinned(overrides) });
|
|
const external = (id = "mcp-a") => ({ sessionId: id, pinned: {} });
|
|
|
|
function fixture(options = {}) {
|
|
const admins = [];
|
|
const services = [];
|
|
const registry = new PostgresConnectionRegistry({
|
|
...options,
|
|
adminFactory: async (connection) => {
|
|
const admin = {
|
|
connection,
|
|
statements: [],
|
|
ended: false,
|
|
async query(sql, values) {
|
|
this.statements.push(sql);
|
|
if (sql.startsWith("SHOW")) return { rows: [{ server_version_num: "170000" }] };
|
|
return { rowCount: 0, rows: [] };
|
|
},
|
|
async end() { this.ended = true; }
|
|
};
|
|
if (options.adminFails) throw Object.assign(new Error("refused"), { code: "ECONNREFUSED" });
|
|
admins.push(admin);
|
|
return admin;
|
|
},
|
|
serviceFactory: (serviceOptions) => {
|
|
const service = {
|
|
options: serviceOptions,
|
|
destroyed: false,
|
|
async verify() { if (options.verifyFails) throw new Error("no"); return { host: serviceOptions.host, port: serviceOptions.port, database: serviceOptions.database, user: serviceOptions.user }; },
|
|
async query(args) { return { rows: [], args }; },
|
|
async execute(args) { return { command: "OK", args }; },
|
|
async listTables(args) { return { tables: [], args }; },
|
|
async describeTable(args) { return { columns: [], args }; },
|
|
destroy() { this.destroyed = true; }
|
|
};
|
|
services.push(service);
|
|
return service;
|
|
}
|
|
});
|
|
return { registry, admins, services };
|
|
}
|
|
|
|
test("binds an internal session to its execution's database and connects as that database's own user", async () => {
|
|
const { registry, admins, services } = fixture();
|
|
assert.deepEqual(await registry.openInternal(internal()), { database: DATABASE });
|
|
|
|
assert.equal(admins[0].connection.user, "ops", "the catalog's user only makes it");
|
|
assert.ok(admins[0].statements.includes(`CREATE DATABASE "${DATABASE}" OWNER "${DATABASE}"`));
|
|
assert.ok(admins[0].ended, "and is let go of");
|
|
assert.equal(services[0].options.database, DATABASE);
|
|
assert.equal(services[0].options.user, DATABASE, "the user has the database's name");
|
|
assert.equal(services[0].options.password, rolePassword(SECRET, DATABASE));
|
|
assert.equal(services[0].options.sslmode, "require", "encrypted unless the catalog says otherwise");
|
|
assert.ok(admins[0].statements.includes(`ALTER ROLE "${DATABASE}" LOGIN PASSWORD '${rolePassword(SECRET, DATABASE)}'`));
|
|
});
|
|
|
|
test("takes the port and the SSL mode from the catalog's headers when it sets them", async () => {
|
|
const { registry, services } = fixture();
|
|
await registry.openInternal(internal("mcp-a", { "x-postgres-port": "5433", "x-postgres-sslmode": "disable" }));
|
|
assert.equal(services[0].options.port, 5433);
|
|
assert.equal(services[0].options.sslmode, "disable");
|
|
});
|
|
|
|
test("prepares an internal session's database once", async () => {
|
|
const { registry, admins } = fixture();
|
|
await registry.openInternal(internal());
|
|
await registry.openInternal(internal());
|
|
assert.equal(admins.length, 1);
|
|
});
|
|
|
|
test("routes the four operations to the session's own connection, and nothing else", async () => {
|
|
const { registry } = fixture();
|
|
const context = internal();
|
|
await registry.openInternal(context);
|
|
assert.deepEqual((await registry.callInternal("query", { statement: "SELECT 1" }, context)).args, { statement: "SELECT 1" });
|
|
assert.equal((await registry.callInternal("execute", { statement: "CREATE TABLE t (id int)" }, context)).command, "OK");
|
|
await registry.callInternal("listTables", {}, context);
|
|
await registry.callInternal("describeTable", { table: "t" }, context);
|
|
await assert.rejects(registry.callInternal("destroy", {}, context), /Unsupported PostgreSQL operation/);
|
|
await assert.rejects(registry.callInternal("query", {}, internal("never-opened")), /no execution storage/);
|
|
});
|
|
|
|
test("refuses a scope that cannot make a database name", async () => {
|
|
const { registry } = fixture();
|
|
for (const scope of ["", "Has Space", "x".repeat(60), "a;drop", "a.b"]) {
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-scope": scope })), /x-postgres-scope/, scope);
|
|
}
|
|
});
|
|
|
|
test("keeps every internal database behind the prefix, whatever the scope names", async () => {
|
|
const { registry, services } = fixture({ internalDatabasePrefix: "hf_run_" });
|
|
await registry.openInternal(internal());
|
|
assert.equal(services[0].options.database, `hf_run_${EXECUTION.replaceAll("-", "_")}`);
|
|
for (const internalDatabasePrefix of ["", "Exec_", "9x", "x".repeat(27)]) {
|
|
assert.throws(() => new PostgresConnectionRegistry({ internalDatabasePrefix }), /prefix/);
|
|
}
|
|
});
|
|
|
|
test("refuses an internal session missing its connection headers, or with a secret too short to trust", async () => {
|
|
const { registry } = fixture();
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-host": undefined })), /needs x-postgres-host/);
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-role-secret": undefined })), /x-postgres-role-secret/);
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-role-secret": "short" })), /at least 32 characters/);
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-host": "http://x" })), /host/);
|
|
await assert.rejects(registry.openInternal(internal("mcp-a", { "x-postgres-port": "99999" })), /port/);
|
|
});
|
|
|
|
test("holds the internal host to the server allowlist", async () => {
|
|
const { registry } = fixture({ allowedHosts: ["db.example.org"] });
|
|
await assert.rejects(registry.openInternal(internal()), /not on the server allowlist/);
|
|
const allowed = fixture({ allowedHosts: ["DB.internal"] });
|
|
await allowed.registry.openInternal(internal());
|
|
});
|
|
|
|
test("says what to fix when the catalog's user cannot make databases, without its address or password", async () => {
|
|
const denying = new PostgresConnectionRegistry({
|
|
adminFactory: async () => ({
|
|
async query(sql) { if (sql.startsWith("CREATE ROLE")) throw Object.assign(new Error("permission denied: ops-password"), { code: "42501" }); return { rowCount: 0, rows: [] }; },
|
|
async end() {}
|
|
})
|
|
});
|
|
await assert.rejects(denying.openInternal(internal()), (error) => /CREATEDB and CREATEROLE/.test(error.message) && !/ops-password|db\.internal/.test(error.message));
|
|
const unreachable = fixture({ adminFails: true });
|
|
await assert.rejects(unreachable.registry.openInternal(internal()), (error) => /Could not reach/.test(error.message) && !/db\.internal/.test(error.message));
|
|
});
|
|
|
|
test("releases the connection of an internal session whose database cannot be had", async () => {
|
|
const { registry, services } = fixture({ verifyFails: true });
|
|
await assert.rejects(registry.openInternal(internal()));
|
|
assert.equal(services[0].destroyed, true);
|
|
});
|
|
|
|
test("keeps internal sessions apart and closes one with its MCP session", async () => {
|
|
const { registry, services } = fixture();
|
|
await registry.openInternal(internal("mcp-a"));
|
|
await registry.openInternal(internal("mcp-b", { "x-postgres-scope": "00000000-0000-0000-0000-000000000002" }));
|
|
assert.notEqual(services[0].options.database, services[1].options.database);
|
|
registry.closeMcpSession("mcp-a");
|
|
assert.equal(services[0].destroyed, true);
|
|
assert.equal(services[1].destroyed, false);
|
|
});
|
|
|
|
const externalOn = (host, port, id = "mcp-a", database = "d") =>
|
|
({ sessionId: id, pinned: { "x-postgres-host": host, "x-postgres-database": database, ...(port === undefined ? {} : { "x-postgres-port": port }) } });
|
|
const publicHost = async (host) => ({ address: host === "db.example.org" ? "203.0.113.10" : "198.51.100.7" });
|
|
|
|
test("an external session goes to the host and port of the node's configuration, and the model gives only the credentials", async () => {
|
|
const { registry, services } = fixture({ resolveHost: publicHost });
|
|
const opened = await registry.open({ user: "u", password: "p" }, externalOn("DB.example.org", "6543"));
|
|
assert.equal(opened.database, "d");
|
|
assert.equal(services[0].options.host, "db.example.org");
|
|
assert.equal(services[0].options.port, 6543);
|
|
assert.equal(services[0].options.connectHost, "203.0.113.10", "the address that was checked is the one connected to");
|
|
assert.equal(services[0].options.sslmode, "require");
|
|
await assert.rejects(registry.open({ host: "elsewhere.example.org", user: "u", password: "p" }, externalOn("db.example.org")),
|
|
/Unknown argument: host/, "the model cannot choose where to connect");
|
|
});
|
|
|
|
test("an empty port from an unset configuration means the default one", async () => {
|
|
const { registry, services } = fixture({ resolveHost: publicHost });
|
|
await registry.open({ user: "u", password: "p" }, externalOn("db.example.org", ""));
|
|
assert.equal(services[0].options.port, 5432);
|
|
});
|
|
|
|
test("an external session without the node's host says where it should come from", async () => {
|
|
const { registry } = fixture({ resolveHost: publicHost });
|
|
await assert.rejects(registry.open({ user: "u", password: "p" }, external()), /comes from the node's configuration, sent as x-postgres-host/);
|
|
await assert.rejects(registry.open({ user: "u", password: "p" }, externalOn("")), /x-postgres-host/);
|
|
await assert.rejects(registry.open({ user: "u", password: "p" }, externalOn("db.example.org", "99999")), /port/);
|
|
await assert.rejects(registry.open({ user: "u", password: "p" }, externalOn("db.example.org", undefined, "mcp-a", "")), /sent as x-postgres-database/);
|
|
});
|
|
|
|
test("without an allowlist only a public address is reached, and with one only the hosts it lists", async () => {
|
|
const refusing = fixture({ resolveHost: async () => { throw new Error("This PostgreSQL host resolves to an address that is not a public one"); } });
|
|
await assert.rejects(refusing.registry.open({ user: "u", password: "p" }, externalOn("internal.example.org")), /not a public one/);
|
|
assert.equal(refusing.services.length, 0, "nothing is opened toward it");
|
|
|
|
const listed = fixture({ allowedHosts: ["db.internal"], resolveHost: async () => { throw new Error("must not be asked"); } });
|
|
await listed.registry.open({ user: "u", password: "p" }, externalOn("db.internal"));
|
|
assert.equal(listed.services[0].options.connectHost, undefined, "a listed host is the operator's choice, private or not");
|
|
await assert.rejects(listed.registry.open({ user: "u", password: "p" }, externalOn("db.example.org")), /not on the server allowlist/);
|
|
});
|
|
|
|
test("keeps external sessions by id to their MCP session, within the limit", async () => {
|
|
const { registry, services } = fixture({ resolveHost: publicHost, maxConnectionsPerSession: 2 });
|
|
const context = externalOn("db.example.org");
|
|
const opened = await registry.open({ user: "u", password: "p" }, context);
|
|
assert.deepEqual((await registry.call("query", { postgresSessionId: opened.postgresSessionId, statement: "SELECT 1" }, context)).args, { statement: "SELECT 1" });
|
|
await assert.rejects(registry.call("query", { postgresSessionId: opened.postgresSessionId, statement: "SELECT 1" }, externalOn("db.example.org", undefined, "mcp-b")), /Unknown PostgreSQL session/);
|
|
await assert.rejects(registry.call("query", { statement: "SELECT 1" }, context), /postgresSessionId is required/);
|
|
|
|
await registry.open({ user: "u", password: "p" }, context);
|
|
await assert.rejects(registry.open({ user: "u", password: "p" }, context), /maximum of 2/);
|
|
|
|
assert.deepEqual(registry.close({ postgresSessionId: opened.postgresSessionId }, context), { postgresSessionId: opened.postgresSessionId, closed: true });
|
|
assert.equal(services[0].destroyed, true);
|
|
});
|
|
|
|
test("a session that asked for the internal storage is told apart from one that did not, even with an empty value", () => {
|
|
assert.equal(internalScopeOf({ pinned: { "x-postgres-scope": "" } }), "");
|
|
assert.equal(internalScopeOf({ pinned: {} }), null);
|
|
assert.ok(PINNED_HEADERS.includes("x-postgres-role-secret"));
|
|
});
|
|
|
|
test("an external session never makes a database or a user: a person's own database is theirs to make", async () => {
|
|
let made = 0;
|
|
const registry = new PostgresConnectionRegistry({
|
|
resolveHost: async () => ({ address: "198.51.100.7" }),
|
|
adminFactory: async () => { made++; throw new Error("must not be asked"); },
|
|
serviceFactory: (options) => ({ options, async verify() { return { database: options.database }; }, destroy() {} })
|
|
});
|
|
await registry.open({ user: "u", password: "p" }, { sessionId: "mcp-a", pinned: { "x-postgres-host": "db.example.org", "x-postgres-database": "mine" } });
|
|
assert.equal(made, 0);
|
|
});
|