PostgreSQL RLS with JWTs for Microservice Tenancy Control

-- Create a sample organizations table
CREATE TABLE organizations (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name TEXT NOT NULL
);

-- Create a products table with a foreign key to organizations
CREATE TABLE products (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    organization_id UUID NOT NULL REFERENCES organizations(id),
    name TEXT NOT NULL,
    price NUMERIC(10, 2) NOT NULL
);

-- Create an index on the tenancy column, as it will be used in every RLS check
CREATE INDEX idx_products_organization_id ON products(organization_id);

-- Insert some sample data for two different tenants
INSERT INTO organizations (id, name) VALUES
    ('a4e3c2b1-3b6a-4b0d-8215-25d2f3a8b4f2', 'Org A'),
    ('b5f4d3c2-4c7b-5c1e-9326-36e3g4b9c5g3', 'Org B');

INSERT INTO products (organization_id, name, price) VALUES
    ('a4e3c2b1-3b6a-4b0d-8215-25d2f3a8b4f2', 'Widget A', 99.99),
    ('b5f4d3c2-4c7b-5c1e-9326-36e3g4b9c5g3', 'Gadget B', 149.50);

-- Enable Row Level Security on the products table
ALTER TABLE products ENABLE ROW LEVEL SECURITY;

-- It's a best practice to also force RLS for the table owner.
-- This prevents accidental bypass during migrations or administrative tasks.
ALTER TABLE products FORCE ROW LEVEL SECURITY;

-- A helper function to get a claim, returning NULL if not found.
-- The second parameter to current_setting('...', true) makes it not throw an error if the setting is missing.
CREATE OR REPLACE FUNCTION get_jwt_claim(claim_name TEXT) RETURNS TEXT AS $$
BEGIN
    RETURN current_setting('request.jwt.claims.' || claim_name, true);
EXCEPTION
    WHEN UNDEFINED_OBJECT THEN
        RETURN NULL;
END;
$$ LANGUAGE plpgsql STABLE;

-- Policy for SELECT operations
-- This policy allows a user to see products if their org_id claim matches the product's organization_id
CREATE POLICY select_products_own_org
ON products
FOR SELECT
USING (organization_id::text = get_jwt_claim('org_id'));

-- Policy for INSERT operations
-- The WITH CHECK clause is used for INSERT and UPDATE.
-- It ensures that any new row being created (or an existing row being moved)
-- adheres to the policy.
CREATE POLICY insert_products_own_org
ON products
FOR INSERT
WITH CHECK (organization_id::text = get_jwt_claim('org_id'));

-- Policy for UPDATE operations
-- 'USING' applies to which rows can be seen and therefore updated.
-- 'WITH CHECK' ensures you can't update a row to move it to an organization you don't belong to.
CREATE POLICY update_products_own_org
ON products
FOR UPDATE
USING (organization_id::text = get_jwt_claim('org_id'))
WITH CHECK (organization_id::text = get_jwt_claim('org_id'));

-- Policy for DELETE operations
-- Only 'USING' is needed here, as you can only delete rows you can see.
CREATE POLICY delete_products_own_org
ON products
FOR DELETE
USING (organization_id::text = get_jwt_claim('org_id'));

// file: jwt-rls-middleware.js
const jwt = require('jsonwebtoken');
const { Pool } = require('pg');

const JWT_SECRET = 'your-super-secret-key'; // Use environment variables in production!

// Assume you have a configured pg Pool instance
const pool = new Pool({
    connectionString: 'postgresql://user:password@host:port/database'
});

// A mock function to generate tokens for testing
function generateTestToken(userId, orgId, role) {
    const payload = {
        sub: userId,
        org_id: orgId,
        role: role,
    };
    return jwt.sign(payload, JWT_SECRET, { expiresIn: '1h' });
}

// The core middleware
async function rlsMiddleware(req, res, next) {
    const authHeader = req.headers['authorization'];
    if (!authHeader || !authHeader.startsWith('Bearer ')) {
        return res.status(401).send('Unauthorized: No token provided.');
    }

    const token = authHeader.substring(7);

    let decodedToken;
    try {
        decodedToken = jwt.verify(token, JWT_SECRET);
    } catch (err) {
        return res.status(401).send('Unauthorized: Invalid token.');
    }

    // Get a client from the pool. This client will be used for the entire request.
    const client = await pool.connect();
    req.dbClient = client; // Attach client to the request object

    try {
        // CRITICAL: Start a transaction.
        await client.query('BEGIN');

        // Securely set the session variables using parameterized queries to prevent SQL injection.
        // We use SET LOCAL so the setting only lasts for the current transaction.
        await client.query(`SET LOCAL "request.jwt.claims.sub" = $1`, [decodedToken.sub]);
        await client.query(`SET LOCAL "request.jwt.claims.org_id" = $1`, [decodedToken.org_id]);
        await client.query(`SET LOCAL "request.jwt.claims.role" = $1`, [decodedToken.role]);

        // Proceed to the actual route handler
        next();

    } catch (err) {
        // If setting claims fails, rollback and release the client
        await client.query('ROLLBACK');
        client.release();
        console.error('Failed to set RLS context:', err);
        res.status(500).send('Internal Server Error');
    }
}

// Middleware to commit/rollback and release the client after the request is handled
async function transactionFinalizer(req, res, next) {
    if (!req.dbClient) return next();

    // The response has been sent, so we can clean up.
    // We listen for the 'finish' event on the response object.
    res.on('finish', async () => {
        try {
            // Check if the transaction is still active before trying to commit/rollback
            // This handles cases where the route handler might have already committed.
            const { rows } = await req.dbClient.query("SELECT txid_current_if_assigned() IS NOT NULL AS in_transaction");
            if (rows[0].in_transaction) {
                // If headers have been sent and status is 2xx, commit.
                // Otherwise, rollback.
                if (res.headersSent && res.statusCode >= 200 && res.statusCode < 300) {
                    await req.dbClient.query('COMMIT');
                } else {
                    await req.dbClient.query('ROLLBACK');
                }
            }
        } catch (err) {
            console.error('Error during transaction finalization:', err);
            // Even if finalization fails, we MUST release the client
        } finally {
            req.dbClient.release();
        }
    });
    next();
}

module.exports = { rlsMiddleware, transactionFinalizer, generateTestToken };

// file: server.js
const express = require('express');
const { rlsMiddleware, transactionFinalizer, generateTestToken } = require('./jwt-rls-middleware');

const app = express();

// Apply middleware to all protected routes
app.use('/api', rlsMiddleware, transactionFinalizer);

app.get('/api/products', async (req, res) => {
    try {
        // The business logic is now beautifully simple.
        // No need to pass org_id or user_id around.
        const { rows } = await req.dbClient.query('SELECT * FROM products');
        res.status(200).json(rows);
    } catch (err) {
        // The finalizer middleware will handle rollback
        console.error(err);
        res.status(500).send('Error fetching products');
    }
});

// Example of how to generate tokens for testing with cURL
app.get('/token', (req, res) => {
    const tokenA = generateTestToken('user-1', 'a4e3c2b1-3b6a-4b0d-8215-25d2f3a8b4f2', 'member');
    const tokenB = generateTestToken('user-2', 'b5f4d3c2-4c7b-5c1e-9326-36e3g4b9c5g3', 'member');
    res.json({ for_org_a: tokenA, for_org_b: tokenB });
});

app.listen(3000, () => console.log('Server running on port 3000'));

-- Add a policy to allow users with the 'admin' role to see all products in their org
-- For a support tool, this could be a global admin role that bypasses the org check entirely
CREATE POLICY select_products_admin_all_orgs
ON products
FOR SELECT
USING (get_jwt_claim('role') = 'super_admin');

-- Policy allowing a system worker to access any product
CREATE POLICY system_worker_access
ON products
FOR ALL -- Applies to SELECT, INSERT, UPDATE, DELETE
USING (get_jwt_claim('role') = 'system_worker')
WITH CHECK (get_jwt_claim('role') = 'system_worker');

Seq Scan on products  (cost=0.00..35.50 rows=1 width=60) (actual time=0.032..0.033 rows=1 loops=1)
  Filter: (organization_id = ((current_setting('request.jwt.claims.org_id'::text, true))::uuid))
  Rows Removed by Filter: 1
Planning Time: 0.115 ms
Execution Time: 0.053 ms