From 289de1c7e6a2f15b8b2d1ab4cafbfba0a0bb5f6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Omar=20Mu=C3=B1oz?= Date: Wed, 21 Jan 2026 09:43:52 -0600 Subject: [PATCH] RLS works --- .../example/demo/RLSConnectionManager.java | 44 ++++++++------ .../com/example/demo/RLSTestController.java | 58 +++++++++++++------ src/main/resources/application.properties | 4 +- 3 files changed, 68 insertions(+), 38 deletions(-) diff --git a/src/main/java/com/example/demo/RLSConnectionManager.java b/src/main/java/com/example/demo/RLSConnectionManager.java index 51a47c0..cb31a25 100644 --- a/src/main/java/com/example/demo/RLSConnectionManager.java +++ b/src/main/java/com/example/demo/RLSConnectionManager.java @@ -1,11 +1,13 @@ package com.example.demo; import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.jdbc.datasource.SingleConnectionDataSource; import org.springframework.stereotype.Component; import javax.sql.DataSource; import java.sql.Connection; import java.sql.SQLException; +import java.sql.Statement; import java.util.function.Function; /** @@ -18,11 +20,9 @@ import java.util.function.Function; public class RLSConnectionManager { private final DataSource dataSource; - private final JdbcTemplate jdbcTemplate; - public RLSConnectionManager(DataSource dataSource, JdbcTemplate jdbcTemplate) { + public RLSConnectionManager(DataSource dataSource) { this.dataSource = dataSource; - this.jdbcTemplate = jdbcTemplate; } /** @@ -43,18 +43,27 @@ public class RLSConnectionManager { // CRITICAL: Set connection to manual commit to ensure atomicity connection.setAutoCommit(false); - // Create a JdbcTemplate bound to THIS specific connection - Connection finalConnection = connection; - JdbcTemplate scopedTemplate = new JdbcTemplate(dataSource) { - // @Override - public Connection getConnection() { - return finalConnection; - } - }; - try { - // Set the RLS context variable on THIS connection - scopedTemplate.execute("SET LOCAL app.current_user_id = " + userId); + // Set the RLS context variable on THIS connection using raw Statement + try (Statement stmt = connection.createStatement()) { + stmt.execute("SET LOCAL app.current_user_id = '" + userId + "'"); + System.out.println("RLS context set for user: " + userId); + } + + // Verify it's set (for debugging) + try (Statement stmt = connection.createStatement()) { + var rs = stmt.executeQuery("SELECT current_setting('app.current_user_id', true)"); + if (rs.next()) { + String value = rs.getString(1); + System.out.println("Verified context value: " + value); + } + } + + // Create a JdbcTemplate bound to THIS specific connection + // Use SingleConnectionDataSource to ensure the template uses this exact connection + SingleConnectionDataSource singleConnectionDataSource = + new SingleConnectionDataSource(connection, true); + JdbcTemplate scopedTemplate = new JdbcTemplate(singleConnectionDataSource); // Execute the operation with the scoped template T result = operation.apply(scopedTemplate); @@ -70,11 +79,10 @@ public class RLSConnectionManager { throw new RuntimeException("Error executing RLS operation", e); } finally { // CRITICAL: Reset the context variable before returning connection to pool - // Using RESET is safer than SET because it works even if SET LOCAL wasn't used - try { - scopedTemplate.execute("RESET app.current_user_id"); + try (Statement stmt = connection.createStatement()) { + stmt.execute("RESET app.current_user_id"); + System.out.println("RLS context reset"); } catch (Exception e) { - // Log but don't throw - we still want to close the connection System.err.println("Warning: Failed to reset RLS context: " + e.getMessage()); } } diff --git a/src/main/java/com/example/demo/RLSTestController.java b/src/main/java/com/example/demo/RLSTestController.java index 2df41d8..50f21eb 100644 --- a/src/main/java/com/example/demo/RLSTestController.java +++ b/src/main/java/com/example/demo/RLSTestController.java @@ -69,17 +69,20 @@ public class RLSTestController { "SELECT current_setting('app.current_user_id', true)", String.class ); - result.put("user1Context", ctx); + result.put("user1Context", ctx.isEmpty() ? "EMPTY" : ctx); return null; }); // Check if context leaked (should be null or empty) - String leakedContext = null; + String leakedContext; try { leakedContext = jdbcTemplate.queryForObject( "SELECT current_setting('app.current_user_id', true)", String.class ); + if (leakedContext == null || leakedContext.isEmpty()) { + leakedContext = "NOT_SET"; + } } catch (Exception e) { // Expected - variable should not be set leakedContext = "NOT_SET"; @@ -92,7 +95,7 @@ public class RLSTestController { "SELECT current_setting('app.current_user_id', true)", String.class ); - result.put("user2Context", ctx); + result.put("user2Context", ctx.isEmpty() ? "EMPTY" : ctx); return null; }); @@ -102,6 +105,9 @@ public class RLSTestController { "SELECT current_setting('app.current_user_id', true)", String.class ); + if (leakedContext == null || leakedContext.isEmpty()) { + leakedContext = "NOT_SET"; + } } catch (Exception e) { leakedContext = "NOT_SET"; } @@ -157,27 +163,43 @@ public class RLSTestController { // Enable RLS jdbcTemplate.execute("ALTER TABLE documents ENABLE ROW LEVEL SECURITY"); + // CRITICAL: Force RLS even for table owner (postgres superuser) + // Without this, RLS policies are bypassed for the table owner + jdbcTemplate.execute("ALTER TABLE documents FORCE ROW LEVEL SECURITY"); + // Create RLS policy + // USING clause: determines which rows are visible (for SELECT) + // WITH CHECK clause: determines which rows can be inserted/updated + // Using NULLIF to handle empty strings from current_setting when variable isn't set jdbcTemplate.execute(""" CREATE POLICY user_documents_policy ON documents FOR ALL - USING (user_id = current_setting('app.current_user_id', true)::bigint) + USING (user_id = NULLIF(current_setting('app.current_user_id', true), '')::bigint) + WITH CHECK (user_id = NULLIF(current_setting('app.current_user_id', true), '')::bigint) """); - // Insert test data - jdbcTemplate.update( - "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", - "User 1 Document", "Private content for user 1", 1L - ); - jdbcTemplate.update( - "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", - "User 2 Document", "Private content for user 2", 2L - ); - jdbcTemplate.update( - "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", - "Another User 1 Doc", "More private content for user 1", 1L - ); + // Insert test data WITH RLS context set + // Now that FORCE RLS is enabled, even our inserts must respect the policy + rlsManager.executeWithRLSContext(1L, scopedTemplate -> { + scopedTemplate.update( + "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", + "User 1 Document", "Private content for user 1", 1L + ); + scopedTemplate.update( + "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", + "Another User 1 Doc", "More private content for user 1", 1L + ); + return null; + }); - return "Database setup complete with RLS enabled"; + rlsManager.executeWithRLSContext(2L, scopedTemplate -> { + scopedTemplate.update( + "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)", + "User 2 Document", "Private content for user 2", 2L + ); + return null; + }); + + return "Database setup complete with RLS enabled and FORCED for table owner"; } } diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties index d04d1d2..60ff06e 100644 --- a/src/main/resources/application.properties +++ b/src/main/resources/application.properties @@ -2,6 +2,6 @@ spring.application.name=demo # PostgreSQL Database Configuration spring.datasource.url=jdbc:postgresql://localhost:5432/rls_test -spring.datasource.username=postgres -spring.datasource.password=postgres +spring.datasource.username=rls_test +spring.datasource.password=rls_test spring.datasource.hikari.maximum-pool-size=10