Postgres RLS with JWTs in a Zero-Trust Microservice Architecture

// services/invoices/repository.js

class InvoiceRepository {
  constructor(dbPool) {
    this.pool = dbPool;
  }

  async findById(invoiceId, tenantId) {
    // Manually adding the tenant_id check is required for every query
    const query = 'SELECT * FROM invoices WHERE id = $1 AND tenant_id = $2';
    const { rows } = await this.pool.query(query, [invoiceId, tenantId]);
    return rows[0];
  }

  async findForUser(userId, tenantId, page, limit) {
    // Another manual check, easy to forget
    const query = 'SELECT * FROM invoices WHERE created_by = $1 AND tenant_id = $2 LIMIT $3 OFFSET $4';
    const offset = (page - 1) * limit;
    const { rows } = await this.pool.query(query, [userId, tenantId, limit, offset]);
    return rows;
  }
  
  // ... dozens more methods, each needing manual tenantId/userId checks
}

-- Enable the pgcrypto extension for UUIDs
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- Create our tables
CREATE TABLE tenants (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name TEXT NOT NULL
);

CREATE TABLE users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    email TEXT NOT NULL UNIQUE,
    tenant_id UUID NOT NULL REFERENCES tenants(id)
);

CREATE TABLE invoices (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    tenant_id UUID NOT NULL REFERENCES tenants(id),
    created_by_user_id UUID NOT NULL REFERENCES users(id),
    amount_cents INT NOT NULL,
    status TEXT NOT NULL DEFAULT 'draft',
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- IMPORTANT: Enable Row-Level Security on the tables that need protection
ALTER TABLE invoices ENABLE ROW LEVEL SECURITY;
-- We can also force it, which protects against the owner bypassing it.
ALTER TABLE invoices FORCE ROW LEVEL SECURITY;

-- Create a role for our application to use. It should NOT be a superuser.
CREATE ROLE application_user LOGIN PASSWORD 'your_secure_password';
GRANT CONNECT ON DATABASE your_database TO application_user;
GRANT USAGE ON SCHEMA public TO application_user;
GRANT SELECT, INSERT, UPDATE, DELETE ON invoices, users, tenants TO application_user;

-- This policy enforces strict tenant isolation.
-- It checks a session variable that we will set from our application.
CREATE POLICY tenant_isolation_policy ON invoices
AS PERMISSIVE FOR ALL
TO application_user
USING (tenant_id = current_setting('jwt.claims.tenant_id')::uuid)
WITH CHECK (tenant_id = current_setting('jwt.claims.tenant_id')::uuid);

// middleware/db-context.js
const { Pool } = require('pg');

const pool = new Pool({ /* your connection config */ });

// This function takes JWT claims and sets them as transaction-local variables
async function setAuthContext(client, claims) {
  // Using 'LOCAL' is CRITICAL. It scopes the setting to the current transaction.
  // This prevents settings from leaking between requests when using a connection pool.
  await client.query(`SET LOCAL jwt.claims.tenant_id = '${claims.tenant_id}';`);
  await client.query(`SET LOCAL jwt.claims.user_id = '${claims.sub}';`);
  await client.query(`SET LOCAL jwt.claims.role = '${claims.role}';`);
}

// The middleware factory
function withDB(handler) {
  return async (req, res, next) => {
    const client = await pool.connect();
    try {
      // req.user is populated by a preceding authentication middleware (e.g., passport-jwt)
      if (!req.user || !req.user.tenant_id) {
        return res.status(401).send('Unauthorized: Missing JWT claims.');
      }

      await client.query('BEGIN');
      await setAuthContext(client, req.user); // req.user holds the decoded JWT payload
      
      // Pass the configured client to the actual request handler
      await handler(req, res, client);

      await client.query('COMMIT');
    } catch (error) {
      await client.query('ROLLBACK');
      // Pass to a centralized error handler
      next(error); 
    } finally {
      client.release();
    }
  };
}

module.exports = { withDB, pool };

// services/invoices/repository.js (The NEW version)

class InvoiceRepository {
  // The client is now passed in per-method, already configured and in a transaction.
  async findById(client, invoiceId) {
    // No tenant_id check! RLS handles it.
    const query = 'SELECT * FROM invoices WHERE id = $1';
    const { rows } = await client.query(query, [invoiceId]);
    return rows[0];
  }

  async create(client, { createdByUserId, amountCents }) {
     // The tenant_id for the WITH CHECK clause will be pulled from the session context
     // by the database itself. We still need to provide it as a column value.
     // We can retrieve it from the context if needed for the INSERT statement.
     const tenantIdResult = await client.query(`SELECT current_setting('jwt.claims.tenant_id')::uuid as tenant_id`);
     const tenantId = tenantIdResult.rows[0].tenant_id;

     const query = 'INSERT INTO invoices (tenant_id, created_by_user_id, amount_cents) VALUES ($1, $2, $3) RETURNING *';
     const { rows } = await client.query(query, [tenantId, createdByUserId, amountCents]);
     return rows[0];
  }
}

// routes/invoices.js
const express = require('express');
const { withDB } = require('../middleware/db-context');
const InvoiceRepository = require('../services/invoices/repository');

const router = express.Router();
const repo = new InvoiceRepository();

// The route handler is wrapped in our middleware
router.get('/:id', withDB(async (req, res, client) => {
  const { id } = req.params;
  
  // The repository method is now blissfully unaware of multi-tenancy.
  const invoice = await repo.findById(client, id);
  
  if (!invoice) {
    // This could mean the invoice doesn't exist OR the user can't see it.
    // From a security perspective, we return the same response.
    return res.status(404).send('Invoice not found.');
  }
  
  res.json(invoice);
}));

module.exports = router;

-- Drop the old policy first
DROP POLICY tenant_isolation_policy ON invoices;

-- Create a new policy that handles multiple roles
CREATE POLICY multi_role_tenancy_policy ON invoices
AS PERMISSIVE FOR ALL
TO application_user
USING (
    -- A 'system_admin' role can see everything
    current_setting('jwt.claims.role') = 'system_admin'
    -- Regular users are still restricted to their tenant
    OR tenant_id = current_setting('jwt.claims.tenant_id')::uuid
)
WITH CHECK (
    -- When writing, you MUST be acting within your own tenant.
    -- We don't allow an admin to write records into another tenant's space via this policy.
    tenant_id = current_setting('jwt.claims.tenant_id')::uuid
);

    CREATE ROLE admin_user LOGIN PASSWORD '...';
    ALTER ROLE admin_user BYPASSRLS;
    -- Grant permissions as needed

// test/test-helper.js
const { pool } = require('../middleware/db-context');

async function runAsUser(claims, testFunction) {
  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    // Manually set the context for the test
    await client.query(`SET LOCAL jwt.claims.tenant_id = '${claims.tenant_id}';`);
    await client.query(`SET LOCAL jwt.claims.user_id = '${claims.sub}';`);
    await client.query(`SET LOCAL jwt.claims.role = '${claims.role}';`);
    
    await testFunction(client);
    
    await client.query('ROLLBACK'); // Always rollback in tests to keep DB clean
  } finally {
    client.release();
  }
}

// Usage in a test file
it('should only return invoices for the correct tenant', async () => {
  const userFromTenantA = { sub: 'user-a', tenant_id: 'tenant-a', role: 'user' };
  
  await runAsUser(userFromTenantA, async (client) => {
    const repo = new InvoiceRepository();
    // This should find invoice_a but not invoice_b
    const invoice = await repo.findById(client, 'invoice_a_id');
    expect(invoice).toBeDefined();

    const missingInvoice = await repo.findById(client, 'invoice_b_id');
    expect(missingInvoice).toBeUndefined();
  });
});

-- In psql, as a superuser to see the plan
SET SESSION AUTHORIZATION application_user;
BEGIN;
SET LOCAL jwt.claims.tenant_id = '...';
EXPLAIN (ANALYZE) SELECT * FROM invoices WHERE id = '...';
ROLLBACK;