Last active
May 7, 2023 16:37
-
-
Save Static-Flow/b0bc489419f204d31e81ab401b526576 to your computer and use it in GitHub Desktop.
SQL security idea: Is this a safe implementation that could allow exposing this database publicly and providing access to users provided they were only given data_role permissions?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| CREATE TABLE IF NOT EXISTS data | |
| ( | |
| pkey bigserial PRIMARY KEY, | |
| value text NOT NULL, | |
| count integer NOT NULL, | |
| first_blood text NOT NULL | |
| ); | |
| CREATE UNIQUE INDEX IF NOT EXISTS value_idx ON data USING btree(value); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| CREATE OR REPLACE FUNCTION insert_value(value_input text) RETURNS void SECURITY DEFINER AS $$ | |
| BEGIN | |
| IF take_token() then | |
| INSERT INTO data (value,first_blood, count) VALUES (value_input, session_user,1) ON CONFLICT (value) DO UPDATE SET count = data.count + 1; | |
| END IF; | |
| END; | |
| $$ LANGUAGE plpgsql; |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| CREATE ROLE data_role; | |
| GRANT EXECUTE ON FUNCTION insert_value(text) TO data_role; | |
| GRANT EXECUTE ON FUNCTION take_token() TO data_role; |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| CREATE OR REPLACE FUNCTION take_token () RETURNS boolean SECURITY DEFINER AS $$ | |
| DECLARE | |
| rate INTEGER; | |
| tokens INTEGER; | |
| extra_tokens INTEGER; | |
| new_tokens INTEGER; | |
| last_refill TIMESTAMP; | |
| this_refill TIMESTAMP; | |
| BEGIN | |
| -- Check if this user exists | |
| SELECT per_hour INTO rate FROM token_rates r WHERE r.username = session_user; | |
| IF rate IS NULL THEN | |
| raise notice 'User % does not have a rate configured', session_user; | |
| RETURN FALSE; | |
| END IF; | |
| -- Lock the buckets until end of transaction | |
| LOCK TABLE token_buckets IN EXCLUSIVE MODE; | |
| -- Read current tokens and last take | |
| SELECT b.tokens, b.last_refill INTO tokens, last_refill FROM token_buckets b WHERE b.username = session_user; | |
| IF tokens IS NULL THEN | |
| tokens := rate; -- Start with the max amount of tokens | |
| last_refill = now(); | |
| raise notice 'Setting up a bucket for user % with % tokens',session_user, tokens; | |
| INSERT INTO token_buckets VALUES (session_user, tokens, last_refill); | |
| END IF; | |
| -- Calculate newly generated tokens since last call | |
| extra_tokens := floor( | |
| EXTRACT(EPOCH FROM (now() - last_refill) * rate / 3600.0) | |
| )::int; | |
| this_refill := last_refill + (extra_tokens * interval '1 second' * 3600.0 / rate); | |
| new_tokens := LEAST(rate, tokens + extra_tokens); | |
| raise notice 'User % has % tokens, last batch generated at %', session_user, new_tokens, this_refill; | |
| -- If there are no tokens left then we don't need to do anything | |
| IF new_tokens <= 0 THEN | |
| RETURN FALSE; | |
| END IF; | |
| -- Set new values and return | |
| UPDATE token_buckets b SET (tokens, last_refill) = (new_tokens - 1, this_refill) WHERE b.username = session_user; | |
| RETURN TRUE; | |
| END | |
| $$ LANGUAGE plpgsql; |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| DROP TABLE IF EXISTS token_buckets; | |
| CREATE TABLE token_buckets ( | |
| username text PRIMARY KEY, | |
| tokens INTEGER, | |
| last_refill TIMESTAMP | |
| ); | |
| DROP TABLE IF EXISTS token_rates; | |
| CREATE TABLE token_rates ( | |
| username text PRIMARY KEY, | |
| per_hour INTEGER | |
| ); |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment