fix: add row level security over shared tables

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;

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;
 

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);

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]);

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()

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()