diff --git a/sql/pg_net.sql b/sql/pg_net.sql index b0acd86a..630ed191 100644 --- a/sql/pg_net.sql +++ b/sql/pg_net.sql @@ -15,7 +15,18 @@ create unlogged table net.http_request_queue( url text not null, headers jsonb, body bytea, - timeout_milliseconds int not null + timeout_milliseconds int not null, + rolename text not null default 'public' +); +alter table net.http_request_queue enable row level security; + +create policy net_request_policy +on net.http_request_queue +using ( + case + when rolename = 'public' then true + else pg_has_role(current_user, rolename, 'member') + end ); create or replace function net.check_worker_is_up() returns void as $$ @@ -40,11 +51,22 @@ create unlogged table net._http_response( content text, timed_out bool, error_msg text, - created timestamptz not null default now() + created timestamptz not null default now(), + rolename text not null ); +alter table net._http_response enable row level security; create index on net._http_response (created); +create policy net_response_policy +on net._http_response +using ( + case + when rolename = 'public' then true + else pg_has_role(current_user, rolename, 'member') + end +); + -- Blocks until an http_request is complete -- API: Private create or replace function net._await_response( @@ -116,7 +138,9 @@ create or replace function net.http_get( -- key/values to be included in request headers headers jsonb default '{}'::jsonb, -- the maximum number of milliseconds the request may take before being cancelled - timeout_milliseconds int default 5000 + timeout_milliseconds int default 5000, + -- the role that should have access to the request/response data + rolename text default 'public' ) -- request_id reference returns bigint @@ -131,12 +155,13 @@ begin from jsonb_each_text(params); -- Add to the request queue - insert into net.http_request_queue(method, url, headers, timeout_milliseconds) + insert into net.http_request_queue(method, url, headers, timeout_milliseconds, rolename) values ( 'GET', net._encode_url_with_params_array(url, params_array), headers, - timeout_milliseconds + timeout_milliseconds, + rolename ) returning id into request_id; @@ -159,7 +184,9 @@ create or replace function net.http_post( -- key/values to be included in request headers headers jsonb default '{"Content-Type": "application/json"}'::jsonb, -- the maximum number of milliseconds the request may take before being cancelled - timeout_milliseconds int DEFAULT 5000 + timeout_milliseconds int DEFAULT 5000, + -- the role that should have access to the request/response data + rolename text default 'public' ) -- request_id reference returns bigint @@ -200,13 +227,14 @@ begin jsonb_each_text(params); -- Add to the request queue - insert into net.http_request_queue(method, url, headers, body, timeout_milliseconds) + insert into net.http_request_queue(method, url, headers, body, timeout_milliseconds, rolename) values ( 'POST', net._encode_url_with_params_array(url, params_array), headers, convert_to(body::text, 'UTF8'), - timeout_milliseconds + timeout_milliseconds, + rolename ) returning id into request_id; @@ -229,7 +257,9 @@ create or replace function net.http_delete( -- the maximum number of milliseconds the request may take before being cancelled timeout_milliseconds int default 5000, -- optional body of the request - body jsonb default NULL + body jsonb default NULL, + -- the role that should have access to the request/response data + rolename text default 'public' ) -- request_id reference returns bigint @@ -244,13 +274,14 @@ begin from jsonb_each_text(params); -- Add to the request queue - insert into net.http_request_queue(method, url, headers, body, timeout_milliseconds) + insert into net.http_request_queue(method, url, headers, body, timeout_milliseconds, rolename) values ( 'DELETE', net._encode_url_with_params_array(url, params_array), headers, convert_to(body::text, 'UTF8'), - timeout_milliseconds + timeout_milliseconds, + rolename ) returning id into request_id; diff --git a/src/core.c b/src/core.c index a7f0824a..0f0c8f94 100644 --- a/src/core.c +++ b/src/core.c @@ -98,7 +98,8 @@ void init_curl_handle(CurlHandle *handle, RequestQueueRow row) { handle->request_headers = request_headers; } - handle->url = TextDatumGetCString(row.url); + handle->url = TextDatumGetCString(row.url); + handle->rolename = TextDatumGetCString(row.rolename); handle->req_body = !row.bodyBin.isnull ? TextDatumGetCString(row.bodyBin.value) : NULL; @@ -219,7 +220,7 @@ uint64 consume_request_queue(const int batch_size) { )\ DELETE FROM net.http_request_queue q\ USING rows WHERE q.id = rows.id\ - RETURNING q.id, q.method, q.url, timeout_milliseconds, array(select key || ': ' || value from jsonb_each_text(q.headers)), q.body", + RETURNING q.id, q.method, q.url, timeout_milliseconds, array(select key || ': ' || value from jsonb_each_text(q.headers)), q.body, q.rolename", 1, (Oid[]){INT4OID}); } @@ -257,7 +258,10 @@ RequestQueueRow get_request_queue_row(HeapTuple spi_tupval, TupleDesc spi_tupdes NullableDatum bodyBin = {.value = SPI_getbinval(spi_tupval, spi_tupdesc, 6, &tupIsNull), .isnull = tupIsNull}; - return (RequestQueueRow){id, method, url, timeout_milliseconds, headersBin, bodyBin}; + Datum rolename = SPI_getbinval(spi_tupval, spi_tupdesc, 7, &tupIsNull); + EREPORT_NULL_ATTR(tupIsNull, rolename); + + return (RequestQueueRow){id, method, url, timeout_milliseconds, headersBin, bodyBin, rolename}; } static Jsonb *jsonb_headers_from_curl_handle(CURL *ez_handle) { @@ -278,15 +282,16 @@ static Jsonb *jsonb_headers_from_curl_handle(CURL *ez_handle) { return PG_JSONB_OBJECT_FINISH(headers); } -enum { response_nparams = 7 }; // using an enum because const size_t doesn't compile +enum { response_nparams = 8 }; // using an enum because const size_t doesn't compile static void execute_insert_response(Datum *vals, char *nulls) { if (ins_response_plan == NULL) { - ins_response_plan = prepare_and_save_plan( - "\ - insert into net._http_response(id, status_code, content, headers, content_type, timed_out, error_msg) values ($1, $2, $3, $4, $5, $6, $7)", - response_nparams, - (Oid[response_nparams]){INT8OID, INT4OID, TEXTOID, JSONBOID, TEXTOID, BOOLOID, TEXTOID}); + ins_response_plan = + prepare_and_save_plan("\ + insert into net._http_response(id, status_code, content, headers, content_type, timed_out, error_msg, rolename) values ($1, $2, $3, $4, $5, $6, $7, $8)", + response_nparams, + (Oid[response_nparams]){INT8OID, INT4OID, TEXTOID, JSONBOID, TEXTOID, + BOOLOID, TEXTOID, TEXTOID}); } int ret_code = SPI_execute_plan(ins_response_plan, vals, nulls, false, 0); @@ -307,6 +312,8 @@ void insert_rejected_response(CurlHandle *handle) { nulls[5] = ' '; vals[6] = CStringGetTextDatum(handle->rejected_reason); nulls[6] = ' '; + vals[7] = CStringGetTextDatum(handle->rolename); + nulls[7] = ' '; execute_insert_response(vals, nulls); } @@ -320,6 +327,9 @@ void insert_response(CurlHandle *handle, CURLcode curl_return_code) { vals[0] = Int64GetDatum(handle->id); nulls[0] = ' '; + vals[7] = CStringGetTextDatum(handle->rolename); + nulls[7] = ' '; + if (curl_return_code == CURLE_OK) { Jsonb *jsonb_headers = jsonb_headers_from_curl_handle(handle->ez_handle); long res_http_status_code = 0; @@ -346,6 +356,7 @@ void insert_response(CurlHandle *handle, CURLcode curl_return_code) { vals[5] = BoolGetDatum(false); nulls[5] = ' '; + } else { bool timed_out = curl_return_code == CURLE_OPERATION_TIMEDOUT; diff --git a/src/core.h b/src/core.h index d5c449ed..32391772 100644 --- a/src/core.h +++ b/src/core.h @@ -31,6 +31,7 @@ typedef struct { int32 timeout_milliseconds; NullableDatum headersBin; NullableDatum bodyBin; + Datum rolename; } RequestQueueRow; // The curl easy handle plus additional data, this acts for both the request and @@ -45,6 +46,7 @@ typedef struct { char *method; CURL *ez_handle; char *rejected_reason; // set when the request must not be sent + char *rolename; // set when the request must not be sent } CurlHandle; uint64 delete_expired_responses(char *ttl, int batch_size); diff --git a/src/worker.c b/src/worker.c index f711c954..b030a89f 100644 --- a/src/worker.c +++ b/src/worker.c @@ -323,15 +323,13 @@ void pg_net_worker(__attribute__((unused)) Datum main_arg) { if (requests_consumed > 0) { CurlHandle *handles = palloc(mul_size(sizeof(CurlHandle), requests_consumed)); - // insert_rejected_response() runs SPI, which resets SPI_tuptable, so keep our own pointer // to the consumed rows SPITupleTable *queue_rows = SPI_tuptable; // initialize curl handles for (size_t j = 0; j < requests_consumed; j++) { - init_curl_handle(&handles[j], - get_request_queue_row(queue_rows->vals[j], queue_rows->tupdesc)); + init_curl_handle(&handles[j], get_request_queue_row(queue_rows->vals[j], queue_rows->tupdesc)); if (handles[j].rejected_reason) { insert_rejected_response(&handles[j]); diff --git a/test/test_engine.py b/test/test_engine.py index af8f1450..349e19dc 100644 --- a/test/test_engine.py +++ b/test/test_engine.py @@ -2,4 +2,4 @@ def test_connect(conn): """Sanity test verifying connection to postgres works""" conn.execute("select 1") - assert [(1, )] == conn.execute("select 1").fetchall() + assert [(1,)] == conn.execute("select 1").fetchall() diff --git a/test/test_rls.py b/test/test_rls.py new file mode 100644 index 00000000..2f63e123 --- /dev/null +++ b/test/test_rls.py @@ -0,0 +1,80 @@ +def test_queue_table_insert_default(conn) -> None: + sql = "CREATE USER test_user;" + conn.execute(sql) + + sql = """ + insert into net.http_request_queue(method, url, headers, timeout_milliseconds) + values ( + 'GET', + 'http://localhost:8080', + '{}'::jsonb, + 5000 + ) + returning id; + """ + id = conn.execute(sql).fetchall()[0][0] + conn.commit() + sql = f"SELECT * FROM net.http_request_queue where id = {id};" + response = conn.execute(sql).fetchall() + assert len(response) == 1 + sql = "set role test_user;" + response = conn.execute(sql) + sql = f"SELECT * FROM net.http_request_queue where id = {id};" + response = conn.execute(sql).fetchall() + assert len(response) == 1 + sql = "RESET ROLE;" + response = conn.execute(sql) + sql = "DROP ROLE test_user;" + conn.execute(sql) + conn.commit() + + +def test_queue_table_insert_non_default(conn) -> None: + sql = "CREATE role test_role;" + conn.execute(sql) + sql = "CREATE USER test_user;" + conn.execute(sql) + sql = "GRANT test_role to test_user;" + conn.execute(sql) + + sql = """ + insert into net.http_request_queue(method, url, headers, timeout_milliseconds, rolename) + values ( + 'GET', + 'http://localhost:8080', + '{}'::jsonb, + 5000, + 'test_role' + ) + returning id; + """ + id = conn.execute(sql).fetchall()[0][0] + conn.commit() + sql = f"SELECT * FROM net.http_request_queue where id = {id};" + response = conn.execute(sql).fetchall() + assert len(response) == 1 + sql = f"set role test_user;" + conn.execute(sql) + sql = f"SELECT * FROM net.http_request_queue where id = {id};" + response = conn.execute(sql).fetchall() + assert len(response) == 1 + + sql = "RESET ROLE" + conn.execute(sql) + + sql = "REVOKE test_role FROM test_user;" + conn.execute(sql) + + sql = "set role test_user;" + conn.execute(sql) + + sql = f"SELECT * FROM net.http_request_queue where id = {id};" + response = conn.execute(sql).fetchall() + assert len(response) == 0 + sql = "RESET ROLE;" + conn.execute(sql) + sql = "DROP ROLE test_user;" + conn.execute(sql) + sql = "DROP ROLE test_role;" + conn.execute(sql) + conn.commit()