PostgreSQL RLS with JWTs for Scalable Multi-Tenant Isolation

-- Enable necessary extensions
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- Schema Definition
CREATE TABLE tenants (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
    email TEXT NOT NULL UNIQUE,
    role TEXT NOT NULL DEFAULT 'member', -- e.g., 'member', 'admin'
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE projects (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
    name TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- Crucial Index for RLS Performance
CREATE INDEX idx_projects_tenant_id ON projects(tenant_id);
CREATE INDEX idx_users_tenant_id ON users(tenant_id);

-- Create a role for our application
CREATE ROLE app_user WITH LOGIN PASSWORD 'a_very_strong_password';

-- Grant basic connection permissions
GRANT CONNECT ON DATABASE your_database_name TO app_user;
GRANT USAGE ON SCHEMA public TO app_user;

-- Grant table-level permissions
GRANT SELECT, INSERT, UPDATE, DELETE ON tenants, users, projects TO app_user;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO app_user;

-- This function sets session variables for the current transaction.
-- The 'is_local' parameter in set_config is TRUE, so the settings only last for the current transaction.
-- This is critical for compatibility with connection poolers.
CREATE OR REPLACE FUNCTION set_auth_context(jwt_claims JSONB)
RETURNS VOID AS $$
DECLARE
    tenant_id_val UUID;
    user_id_val UUID;
    role_val TEXT;
BEGIN
    -- Extract claims, ensuring they exist and are of the correct type.
    -- The '->>' operator extracts a JSON field as text.
    tenant_id_val := (jwt_claims ->> 'tenant_id')::UUID;
    user_id_val := (jwt_claims ->> 'user_id')::UUID;
    role_val := (jwt_claims ->> 'role')::TEXT;

    IF tenant_id_val IS NULL OR user_id_val IS NULL OR role_val IS NULL THEN
        RAISE EXCEPTION 'Invalid JWT claims: one or more required fields are missing';
    END IF;

    -- Set the configuration variables for the current transaction.
    -- These are namespaced to avoid conflicts.
    PERFORM set_config('app.current.tenant_id', tenant_id_val::TEXT, true);
    PERFORM set_config('app.current.user_id', user_id_val::TEXT, true);
    PERFORM set_config('app.current.role', role_val, true);

EXCEPTION
    WHEN OTHERS THEN
        -- Log the error and re-raise to ensure the transaction fails.
        RAISE WARNING 'Error setting auth context: %', SQLERRM;
        RAISE;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Grant execute permission to our application user
GRANT EXECUTE ON FUNCTION set_auth_context(JSONB) TO app_user;

-- Enable RLS on the tables we want to protect.
ALTER TABLE projects ENABLE ROW LEVEL SECURITY;
ALTER TABLE users ENABLE ROW LEVEL SECURITY;

-- Create a helper function to get the current tenant_id. This is cleaner than repeating the cast.
-- Mark it as STABLE as its result is stable within a transaction.
CREATE OR REPLACE FUNCTION current_tenant_id()
RETURNS UUID AS $$
BEGIN
    RETURN current_setting('app.current.tenant_id')::UUID;
EXCEPTION
    WHEN OTHERS THEN
        -- Return a default or NULL if the setting is not found.
        -- This prevents queries from failing if run outside an authenticated session.
        -- A policy can then decide how to handle this (e.g., deny access).
        RETURN '00000000-0000-0000-0000-000000000000';
END;
$$ LANGUAGE plpgsql STABLE;

-- **POLICY 1: Simple Tenant Isolation for 'projects' table **
CREATE POLICY tenant_isolation_policy ON projects
AS PERMISSIVE -- PERMISSIVE means policies are combined with OR
FOR ALL -- Applies to SELECT, INSERT, UPDATE, DELETE
TO app_user
USING (tenant_id = current_tenant_id())
WITH CHECK (tenant_id = current_tenant_id());

-- **POLICY 2: More complex policy for 'users' table **
-- Users should only see other users in their own tenant.
-- Additionally, only 'admin' users can see other users. 'member' users can only see themselves.
CREATE POLICY user_visibility_policy ON users
AS PERMISSIVE
FOR SELECT
TO app_user
USING (
    tenant_id = current_tenant_id() AND
    (
        -- Admins can see all users in their tenant
        current_setting('app.current.role') = 'admin' OR
        -- Members can only see their own record
        id = current_setting('app.current.user_id')::UUID
    )
);

-- Policy for INSERT/UPDATE/DELETE on 'users' table
-- Allow admins to modify users in their tenant, but members can only modify themselves.
CREATE POLICY user_modification_policy ON users
AS PERMISSIVE
FOR INSERT, UPDATE, DELETE
TO app_user
USING (
    tenant_id = current_tenant_id() AND
    (
        current_setting('app.current.role') = 'admin' OR
        id = current_setting('app.current.user_id')::UUID
    )
)
WITH CHECK (
    tenant_id = current_tenant_id() AND
    (
        current_setting('app.current.role') = 'admin' OR
        id = current_setting('app.current.user_id')::UUID
    )
);

const { Pool } = require('pg');

const pool = new Pool({
    user: 'app_user',
    host: 'localhost',
    database: 'your_database_name',
    password: 'a_very_strong_password',
    port: 5432,
});

// A mock JWT verification function
function verifyAndDecodeJwt(token) {
    // In a real app, use a library like 'jsonwebtoken' to verify the signature
    // against your public key or secret.
    if (!token) throw new Error('No token provided');
    // For this example, we'll just decode a mock base64 payload.
    // Real JWT: { "tenant_id": "...", "user_id": "...", "role": "admin" }
    try {
        const payload = JSON.parse(Buffer.from(token.split('.')[1], 'base64').toString());
        if (!payload.tenant_id || !payload.user_id || !payload.role) {
            throw new Error('Invalid JWT claims');
        }
        return payload;
    } catch (e) {
        throw new Error('Invalid JWT format');
    }
}

/**
 * Higher-order function to wrap database operations in a transaction
 * with RLS context set from a JWT.
 * @param {object} jwtClaims - The decoded JWT payload.
 * @param {Function} callback - An async function that receives the client and performs db operations.
 */
async function withRLSTransaction(jwtClaims, callback) {
    const client = await pool.connect();
    try {
        await client.query('BEGIN');

        // Set the RLS context for this transaction
        const claimsJson = JSON.stringify(jwtClaims);
        await client.query('SELECT set_auth_context($1::jsonb)', [claimsJson]);

        // Execute the business logic
        const result = await callback(client);

        await client.query('COMMIT');
        return result;
    } catch (e) {
        await client.query('ROLLBACK');
        console.error('Transaction failed:', e.message);
        throw e;
    } finally {
        client.release();
    }
}

// --- Example Usage ---

async function getProjectsForUser(req) {
    try {
        // 1. Get token from request header
        const token = req.headers.authorization.split(' ')[1];

        // 2. Verify and decode JWT
        const claims = verifyAndDecodeJwt(token);

        // 3. Use the transactional wrapper
        const projects = await withRLSTransaction(claims, async (client) => {
            // 4. The query is simple and tenancy-unaware.
            //    RLS handles the security automatically.
            console.log(`Executing query for tenant: ${claims.tenant_id}`);
            const res = await client.query('SELECT id, name FROM projects');
            return res.rows;
        });

        return { status: 200, body: projects };

    } catch (error) {
        // Handle auth errors, db errors, etc.
        return { status: 401, body: { error: error.message } };
    }
}

// --- Mocking a request and running it ---

async function main() {
    // Mock JWT for Tenant 1, User A (Admin)
    const adminJwtTenant1 = `header.${Buffer.from(JSON.stringify({ tenant_id: 'a4e3c23e-3e5b-4ba5-8274-32b4d36a3d61', user_id: 'd9b7f5b3-3e1a-4e8c-8b8f-3e8a6c4e2f8c', role: 'admin' })).toString('base64')}.signature`;

    // Mock JWT for Tenant 2, User B (Member)
    const memberJwtTenant2 = `header.${Buffer.from(JSON.stringify({ tenant_id: 'b5f4d34f-4f6c-5cb6-9385-43c5e47b4e72', user_id: 'e1c8g6c4-4f2b-5f9d-9c9g-4f9b7d5f3g9d', role: 'member' })).toString('base64')}.signature`;

    console.log('--- Fetching projects as Admin from Tenant 1 ---');
    const req1 = { headers: { authorization: `Bearer ${adminJwtTenant1}` } };
    const res1 = await getProjectsForUser(req1);
    console.log('Result:', res1.body);

    console.log('\n--- Fetching projects as Member from Tenant 2 ---');
    const req2 = { headers: { authorization: `Bearer ${memberJwtTenant2}` } };
    const res2 = await getProjectsForUser(req2);
    console.log('Result:', res2.body);
}

// main(); // You would run this after populating some data.

-- Re-creating our helper function with LEAKPROOF and other optimizations
CREATE OR REPLACE FUNCTION current_tenant_id()
RETURNS UUID AS $$
BEGIN
    -- Using coalesce is slightly cleaner than a full exception block for this case
    RETURN coalesce(current_setting('app.current.tenant_id', true), '00000000-0000-0000-0000-000000000000')::UUID;
END;
$$ LANGUAGE plpgsql STABLE LEAKPROOF;

-- In a psql session, simulate an application request
BEGIN;
-- Set the context
SELECT set_auth_context('{"tenant_id": "a4e3c23e-3e5b-4ba5-8274-32b4d36a3d61", "user_id": "...", "role": "admin"}'::jsonb);

-- Analyze the query
EXPLAIN ANALYZE SELECT * FROM projects WHERE name LIKE 'Q1%';

COMMIT;

CREATE ROLE internal_worker WITH LOGIN PASSWORD '...';
ALTER ROLE internal_worker BYPASSRLS;

-- Modified policy
CREATE POLICY tenant_isolation_policy ON projects
FOR ALL TO app_user
USING (
    current_setting('app.rls.bypass', true) = 'true' OR
    tenant_id = current_tenant_id()
)
WITH CHECK (
    current_setting('app.rls.bypass', true) = 'true' OR
    tenant_id = current_tenant_id()
);

ALTER TABLE tenants ADD COLUMN parent_id UUID REFERENCES tenants(id);

CREATE OR REPLACE FUNCTION get_tenant_descendants(root_tenant_id UUID)
RETURNS TABLE(id UUID) AS $$
BEGIN
    RETURN QUERY
    WITH RECURSIVE tenant_tree AS (
        SELECT t.id FROM tenants t WHERE t.id = root_tenant_id
        UNION ALL
        SELECT t.id FROM tenants t
        INNER JOIN tenant_tree tt ON t.parent_id = tt.id
    )
    SELECT tt.id FROM tenant_tree tt;
END;
$$ LANGUAGE plpgsql STABLE LEAKPROOF;

DROP POLICY tenant_isolation_policy ON projects;

CREATE POLICY hierarchical_tenant_policy ON projects
FOR ALL TO app_user
USING (tenant_id IN (SELECT id FROM get_tenant_descendants(current_tenant_id())))
WITH CHECK (tenant_id = current_tenant_id()); -- Still only allow writing to your own tenant