From 701191d5c5c6e04a339d31191881df3e1ca7a4eb Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 23 Oct 2024 17:50:46 +0300 Subject: [PATCH] feat: add role-based permissions for secrets --- projects/extension/ai/secrets.py | 30 ++++++ projects/extension/sql/ai--0.4.0.sql | 67 ++++++++++++- .../extension/sql/idempotent/014-secrets.sql | 18 ++++ .../sql/idempotent/999-privileges.sql | 12 ++- .../incremental/002-secret_permissions.sql | 9 ++ .../extension/tests/contents/output.expected | 35 ++++++- .../tests/privileges/function.expected | 10 +- .../extension/tests/privileges/table.expected | 18 +++- .../tests/privileges/test_privileges.py | 96 +++++++++++++++++++ .../extension/tests/privileges/view.expected | 26 ++--- projects/extension/tests/vectorizer/server.py | 2 +- scripts/vectorizer-load-test/load.py | 2 +- 12 files changed, 303 insertions(+), 22 deletions(-) create mode 100644 projects/extension/sql/incremental/002-secret_permissions.sql diff --git a/projects/extension/ai/secrets.py b/projects/extension/ai/secrets.py index aee57af2c..fc67987e1 100644 --- a/projects/extension/ai/secrets.py +++ b/projects/extension/ai/secrets.py @@ -20,6 +20,32 @@ def get_guc_value(plpy, setting: str, default: str) -> str: return val +def check_secret_permissions(plpy, secret_name: str) -> bool: + # check if the user has access to all secrets + plan = plpy.prepare( + """ + SELECT 1 + FROM ai.secret_permissions + WHERE name = '*'""", + [], + ) + result = plan.execute([], 1) + if len(result) > 0: + return True + + # check if the user has access to the specific secret + plan = plpy.prepare( + """ + SELECT 1 + FROM ai.secret_permissions + WHERE name = $1 + """, + ["text"], + ) + result = plan.execute([secret_name], 1) + return len(result) > 0 + + def resolve_secret(plpy, secret_name: str) -> str: # first try the guc, then the secrets manager, then error secret_name_lower = secret_name.lower() @@ -45,6 +71,10 @@ def reveal_secret(plpy, secret_name: str) -> str | None: plpy.error("secrets manager is not enabled") return None + if not check_secret_permissions(plpy, secret_name): + plpy.error(f"user does not have access to secret '{secret_name}'") + return None + the_url = urljoin( get_guc_value(plpy, GUC_SECRETS_MANAGER_URL, ""), DEFAULT_SECRETS_MANAGER_PATH, diff --git a/projects/extension/sql/ai--0.4.0.sql b/projects/extension/sql/ai--0.4.0.sql index 533ea0586..b9b965cc8 100644 --- a/projects/extension/sql/ai--0.4.0.sql +++ b/projects/extension/sql/ai--0.4.0.sql @@ -114,6 +114,43 @@ end; $outer_migration_block$; +------------------------------------------------------------------------------- +-- 002-secret_permissions.sql +do $outer_migration_block$ /*002-secret_permissions.sql*/ +declare + _sql text; + _migration record; + _migration_name text = $migration_name$002-secret_permissions.sql$migration_name$; + _migration_body text = +$migration_body$ +create table ai._secret_permissions +( + name text not null check(name = '*' or name ~ '^[A-Za-z0-9_.]+$') +, "role" text not null +, primary key (name, "role") +); +perform pg_catalog.pg_extension_config_dump('ai._secret_permissions'::pg_catalog.regclass, ''); +--only admins will have access to this table +revoke all on ai._secret_permissions from public; + +$migration_body$; +begin + select * into _migration from ai.migration where "name" operator(pg_catalog.=) _migration_name; + if _migration is not null then + raise notice 'migration %s already applied. skipping.', _migration_name; + if _migration.body operator(pg_catalog.!=) _migration_body then + raise warning 'the contents of migration "%s" have changed', _migration_name; + end if; + return; + end if; + _sql = pg_catalog.format(E'do /*%s*/ $migration_body$\nbegin\n%s\nend;\n$migration_body$;', _migration_name, _migration_body); + execute _sql; + insert into ai.migration ("name", body, applied_at_version) + values (_migration_name, _migration_body, $version$0.4.0$version$); +end; +$outer_migration_block$; + + -------------------------------------------------------------------------------- -- 001-openai.sql @@ -3219,6 +3256,24 @@ language plpython3u volatile security invoker set search_path to pg_catalog, pg_temp; +create or replace view ai.secret_permissions as +SELECT * +FROM ai._secret_permissions +WHERE to_regrole("role") is not null AND pg_has_role(current_user, "role", 'member'); + +create or replace function ai.grant_secret(secret_name text, grant_to_role text) returns void +as $func$ + insert into ai._secret_permissions (name, "role") VALUES (secret_name, grant_to_role); +$func$ language sql volatile security invoker +set search_path to pg_catalog, pg_temp; + +create or replace function ai.revoke_secret(secret_name text, revoke_from_role text) returns void +as $func$ + delete from ai._secret_permissions where name = secret_name and "role" = revoke_from_role; +$func$ language sql volatile security invoker +set search_path to pg_catalog, pg_temp; + + -------------------------------------------------------------------------------- -- 999-privileges.sql @@ -3275,7 +3330,8 @@ begin and k.relkind in ('r', 'p', 'S', 'v') -- tables, sequences, and views and (admin, n.nspname, k.relname) not in ( - (false, 'ai', 'migration') -- only admins get any access to this table + (false, 'ai', 'migration'), -- only admins get any access to this table + (false, 'ai', '_secret_permissions') -- only admins get any access to this table ) order by n.nspname, k.relname ) @@ -3309,7 +3365,7 @@ begin and e.extname operator(pg_catalog.=) 'ai' and k.prokind in ('f', 'p') and case - when k.proname operator(pg_catalog.=) 'grant_ai_usage' then admin -- only admins get this function + when k.proname in ('grant_ai_usage', 'grant_secret', 'revoke_secret') then admin -- only admins get these function else true end ) @@ -3317,6 +3373,12 @@ begin raise debug '%', _sql; execute _sql; end loop; + + -- secret permissions + if admin then + -- grant access to all secrets to admin users + insert into ai.secret_permissions (name, "role") VALUES ('*', to_user); + end if; end $func$ language plpgsql volatile security invoker -- gotta have privs to give privs @@ -3396,4 +3458,3 @@ $func$; - diff --git a/projects/extension/sql/idempotent/014-secrets.sql b/projects/extension/sql/idempotent/014-secrets.sql index 402d47c71..3a2915b97 100644 --- a/projects/extension/sql/idempotent/014-secrets.sql +++ b/projects/extension/sql/idempotent/014-secrets.sql @@ -8,3 +8,21 @@ as $python$ $python$ language plpython3u volatile security invoker set search_path to pg_catalog, pg_temp; + + +create or replace view ai.secret_permissions as +SELECT * +FROM ai._secret_permissions +WHERE to_regrole("role") is not null AND pg_has_role(current_user, "role", 'member'); + +create or replace function ai.grant_secret(secret_name text, grant_to_role text) returns void +as $func$ + insert into ai._secret_permissions (name, "role") VALUES (secret_name, grant_to_role); +$func$ language sql volatile security invoker +set search_path to pg_catalog, pg_temp; + +create or replace function ai.revoke_secret(secret_name text, revoke_from_role text) returns void +as $func$ + delete from ai._secret_permissions where name = secret_name and "role" = revoke_from_role; +$func$ language sql volatile security invoker +set search_path to pg_catalog, pg_temp; diff --git a/projects/extension/sql/idempotent/999-privileges.sql b/projects/extension/sql/idempotent/999-privileges.sql index e9693c312..2311405c6 100644 --- a/projects/extension/sql/idempotent/999-privileges.sql +++ b/projects/extension/sql/idempotent/999-privileges.sql @@ -51,7 +51,8 @@ begin and k.relkind in ('r', 'p', 'S', 'v') -- tables, sequences, and views and (admin, n.nspname, k.relname) not in ( - (false, 'ai', 'migration') -- only admins get any access to this table + (false, 'ai', 'migration'), -- only admins get any access to this table + (false, 'ai', '_secret_permissions') -- only admins get any access to this table ) order by n.nspname, k.relname ) @@ -85,7 +86,7 @@ begin and e.extname operator(pg_catalog.=) 'ai' and k.prokind in ('f', 'p') and case - when k.proname operator(pg_catalog.=) 'grant_ai_usage' then admin -- only admins get this function + when k.proname in ('grant_ai_usage', 'grant_secret', 'revoke_secret') then admin -- only admins get these function else true end ) @@ -93,6 +94,12 @@ begin raise debug '%', _sql; execute _sql; end loop; + + -- secret permissions + if admin then + -- grant access to all secrets to admin users + insert into ai.secret_permissions (name, "role") VALUES ('*', to_user); + end if; end $func$ language plpgsql volatile security invoker -- gotta have privs to give privs @@ -169,4 +176,3 @@ begin end loop; end $func$; - diff --git a/projects/extension/sql/incremental/002-secret_permissions.sql b/projects/extension/sql/incremental/002-secret_permissions.sql new file mode 100644 index 000000000..f5a448ead --- /dev/null +++ b/projects/extension/sql/incremental/002-secret_permissions.sql @@ -0,0 +1,9 @@ +create table ai._secret_permissions +( + name text not null check(name = '*' or name ~ '^[A-Za-z0-9_.]+$') +, "role" text not null +, primary key (name, "role") +); +perform pg_catalog.pg_extension_config_dump('ai._secret_permissions'::pg_catalog.regclass, ''); +--only admins will have access to this table +revoke all on ai._secret_permissions from public; diff --git a/projects/extension/tests/contents/output.expected b/projects/extension/tests/contents/output.expected index 6f7cc37cb..a92115966 100644 --- a/projects/extension/tests/contents/output.expected +++ b/projects/extension/tests/contents/output.expected @@ -24,6 +24,7 @@ CREATE EXTENSION function ai.execute_vectorizer(integer) function ai.formatting_python_template(text) function ai.grant_ai_usage(name,boolean) + function ai.grant_secret(text,text) function ai.grant_to() function ai.grant_to(name[]) function ai.indexing_default() @@ -48,6 +49,7 @@ CREATE EXTENSION function ai._resolve_indexing_default() function ai._resolve_scheduling_default() function ai.reveal_secret(text) + function ai.revoke_secret(text,text) function ai.scheduling_default() function ai.scheduling_none() function ai.scheduling_timescaledb(interval,timestamp with time zone,boolean,text) @@ -75,10 +77,30 @@ CREATE EXTENSION function ai._vectorizer_vector_index_exists(name,name,jsonb) sequence ai.vectorizer_id_seq table ai.migration + table ai._secret_permissions table ai.vectorizer table ai.vectorizer_errors + view ai.secret_permissions view ai.vectorizer_status -(74 rows) +(78 rows) + + Table "ai._secret_permissions" + Column | Type | Collation | Nullable | Default | Storage | Compression | Stats target | Description +--------+------+-----------+----------+---------+----------+-------------+--------------+------------- + name | text | | not null | | extended | | | + role | text | | not null | | extended | | | +Indexes: + "_secret_permissions_pkey" PRIMARY KEY, btree (name, role) +Check constraints: + "_secret_permissions_name_check" CHECK (name = '*'::text OR name ~ '^[A-Za-z0-9_.]+$'::text) +Access method: heap + + Index "ai._secret_permissions_pkey" + Column | Type | Key? | Definition | Storage | Stats target +--------+------+------+------------+----------+-------------- + name | text | yes | name | extended | + role | text | yes | role | extended | +primary key, btree, for table "ai._secret_permissions" Table "ai.migration" Column | Type | Collation | Nullable | Default | Storage | Compression | Stats target | Description @@ -97,6 +119,17 @@ Access method: heap name | text | yes | name | extended | primary key, btree, for table "ai.migration" + View "ai.secret_permissions" + Column | Type | Collation | Nullable | Default | Storage | Description +--------+------+-----------+----------+---------+----------+------------- + name | text | | | | extended | + role | text | | | | extended | +View definition: + SELECT name, + role + FROM ai._secret_permissions + WHERE to_regrole(role) IS NOT NULL AND pg_has_role(CURRENT_USER, role::name, 'member'::text); + Table "ai.vectorizer" Column | Type | Collation | Nullable | Default | Storage | Compression | Stats target | Description ---------------+---------+-----------+----------+----------------------------------+----------+-------------+--------------+------------- diff --git a/projects/extension/tests/privileges/function.expected b/projects/extension/tests/privileges/function.expected index ac8045853..0ee3c7d4a 100644 --- a/projects/extension/tests/privileges/function.expected +++ b/projects/extension/tests/privileges/function.expected @@ -176,6 +176,10 @@ f | bob | execute | no | ai | grant_ai_usage(to_user name, admin boolean) f | fred | execute | no | ai | grant_ai_usage(to_user name, admin boolean) f | jill | execute | no | ai | grant_ai_usage(to_user name, admin boolean) + f | alice | execute | YES | ai | grant_secret(secret_name text, grant_to_role text) + f | bob | execute | no | ai | grant_secret(secret_name text, grant_to_role text) + f | fred | execute | no | ai | grant_secret(secret_name text, grant_to_role text) + f | jill | execute | no | ai | grant_secret(secret_name text, grant_to_role text) f | alice | execute | YES | ai | grant_to() f | bob | execute | no | ai | grant_to() f | fred | execute | no | ai | grant_to() @@ -264,6 +268,10 @@ f | bob | execute | no | ai | reveal_secret(secret_name text) f | fred | execute | no | ai | reveal_secret(secret_name text) f | jill | execute | YES | ai | reveal_secret(secret_name text) + f | alice | execute | YES | ai | revoke_secret(secret_name text, revoke_from_role text) + f | bob | execute | no | ai | revoke_secret(secret_name text, revoke_from_role text) + f | fred | execute | no | ai | revoke_secret(secret_name text, revoke_from_role text) + f | jill | execute | no | ai | revoke_secret(secret_name text, revoke_from_role text) f | alice | execute | YES | ai | scheduling_default() f | bob | execute | no | ai | scheduling_default() f | fred | execute | no | ai | scheduling_default() @@ -280,5 +288,5 @@ f | bob | execute | no | ai | vectorizer_queue_pending(vectorizer_id integer) f | fred | execute | no | ai | vectorizer_queue_pending(vectorizer_id integer) f | jill | execute | YES | ai | vectorizer_queue_pending(vectorizer_id integer) -(280 rows) +(288 rows) diff --git a/projects/extension/tests/privileges/table.expected b/projects/extension/tests/privileges/table.expected index 8a244c0fb..aa5ef729b 100644 --- a/projects/extension/tests/privileges/table.expected +++ b/projects/extension/tests/privileges/table.expected @@ -1,5 +1,21 @@ schema | table | user | privilege | granted --------+----------------------+-------+-----------+--------- + ai | _secret_permissions | alice | delete | YES + ai | _secret_permissions | alice | insert | YES + ai | _secret_permissions | alice | select | YES + ai | _secret_permissions | alice | update | YES + ai | _secret_permissions | bob | delete | no + ai | _secret_permissions | bob | insert | no + ai | _secret_permissions | bob | select | no + ai | _secret_permissions | bob | update | no + ai | _secret_permissions | fred | delete | no + ai | _secret_permissions | fred | insert | no + ai | _secret_permissions | fred | select | no + ai | _secret_permissions | fred | update | no + ai | _secret_permissions | jill | delete | no + ai | _secret_permissions | jill | insert | no + ai | _secret_permissions | jill | select | no + ai | _secret_permissions | jill | update | no ai | _vectorizer_q_1 | alice | delete | YES ai | _vectorizer_q_1 | alice | insert | YES ai | _vectorizer_q_1 | alice | select | YES @@ -96,5 +112,5 @@ wiki | post_embedding_store | jill | insert | YES wiki | post_embedding_store | jill | select | YES wiki | post_embedding_store | jill | update | YES -(96 rows) +(112 rows) diff --git a/projects/extension/tests/privileges/test_privileges.py b/projects/extension/tests/privileges/test_privileges.py index 727ced178..b8ff0a7c5 100644 --- a/projects/extension/tests/privileges/test_privileges.py +++ b/projects/extension/tests/privileges/test_privileges.py @@ -2,6 +2,7 @@ import subprocess from pathlib import Path +import psycopg import pytest # skip tests in this module if disabled @@ -84,3 +85,98 @@ def test_function_privileges(): def test_jill_privileges(): psql_file("jill", "privs", "jill.sql") + + +def test_secret_privileges(): + # jill cannot access any secrets + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + got_error = False + try: + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + except Exception: + got_error = True + assert got_error + + # alice can access all the secrets and grant them to jill + with psycopg.connect(db_url("alice", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + cur.execute("select ai.grant_secret('OPENAI_API_KEY', 'jill')") + + # jill can access the secret granted to her but not the other one + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + got_error = False + try: + cur.execute("select ai.reveal_secret('OPENAI_API_KEY_2')") + except Exception: + got_error = True + assert got_error + + # alice can revoke the secret from jill + with psycopg.connect(db_url("alice", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.revoke_secret('OPENAI_API_KEY', 'jill')") + + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + got_error = False + try: + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + except Exception: + got_error = True + assert got_error + + # alice can grant the secret to all keys for jill + with psycopg.connect(db_url("alice", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.grant_secret('*', 'jill')") + + # jill can access all the secrets + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + cur.execute("select ai.reveal_secret('OPENAI_API_KEY_2')") + + # alice can revoke the * privilege from jill + with psycopg.connect(db_url("alice", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + cur.execute("select ai.revoke_secret('*', 'jill')") + + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + got_error = False + try: + cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + except Exception: + got_error = True + assert got_error diff --git a/projects/extension/tests/privileges/view.expected b/projects/extension/tests/privileges/view.expected index c333ec898..d2445e24a 100644 --- a/projects/extension/tests/privileges/view.expected +++ b/projects/extension/tests/privileges/view.expected @@ -1,12 +1,16 @@ - schema | view | user | privilege | granted ---------+-------------------+-------+-----------+--------- - ai | vectorizer_status | alice | select | YES - ai | vectorizer_status | bob | select | no - ai | vectorizer_status | fred | select | no - ai | vectorizer_status | jill | select | YES - wiki | post_embedding | alice | select | YES - wiki | post_embedding | bob | select | no - wiki | post_embedding | fred | select | YES - wiki | post_embedding | jill | select | YES -(8 rows) + schema | view | user | privilege | granted +--------+--------------------+-------+-----------+--------- + ai | secret_permissions | alice | select | YES + ai | secret_permissions | bob | select | no + ai | secret_permissions | fred | select | no + ai | secret_permissions | jill | select | YES + ai | vectorizer_status | alice | select | YES + ai | vectorizer_status | bob | select | no + ai | vectorizer_status | fred | select | no + ai | vectorizer_status | jill | select | YES + wiki | post_embedding | alice | select | YES + wiki | post_embedding | bob | select | no + wiki | post_embedding | fred | select | YES + wiki | post_embedding | jill | select | YES +(12 rows) diff --git a/projects/extension/tests/vectorizer/server.py b/projects/extension/tests/vectorizer/server.py index 2b2efd0d2..815e3bc1a 100644 --- a/projects/extension/tests/vectorizer/server.py +++ b/projects/extension/tests/vectorizer/server.py @@ -179,7 +179,7 @@ async def get_secrets(secret_name: str = Header(None, alias="Secret-Name")): ) # For now, we'll just return the test key if the secret_name matches - if secret_name == "OPENAI_API_KEY": + if secret_name == "OPENAI_API_KEY" or secret_name == "OPENAI_API_KEY_2": return {secret_name: "test"} elif secret_name == "ERROR_SECRET": return JSONResponse( diff --git a/scripts/vectorizer-load-test/load.py b/scripts/vectorizer-load-test/load.py index 272c168b9..f16ebb8cc 100644 --- a/scripts/vectorizer-load-test/load.py +++ b/scripts/vectorizer-load-test/load.py @@ -14,7 +14,7 @@ def load(): print("fetching dataset...") - data = load_dataset(f"Cohere/wikipedia-22-12", 'en', split='train', streaming=True, trust_remote_code=True) + data = load_dataset("Cohere/wikipedia-22-12", 'en', split='train', streaming=True, trust_remote_code=True) to_load = -1 ans = 'q'