RLS works

This commit is contained in:
Omar Muñoz
2026-01-21 09:43:52 -06:00
parent 095799f950
commit 289de1c7e6
3 changed files with 68 additions and 38 deletions

View File

@@ -1,11 +1,13 @@
package com.example.demo; package com.example.demo;
import org.springframework.jdbc.core.JdbcTemplate; import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.SingleConnectionDataSource;
import org.springframework.stereotype.Component; import org.springframework.stereotype.Component;
import javax.sql.DataSource; import javax.sql.DataSource;
import java.sql.Connection; import java.sql.Connection;
import java.sql.SQLException; import java.sql.SQLException;
import java.sql.Statement;
import java.util.function.Function; import java.util.function.Function;
/** /**
@@ -18,11 +20,9 @@ import java.util.function.Function;
public class RLSConnectionManager { public class RLSConnectionManager {
private final DataSource dataSource; private final DataSource dataSource;
private final JdbcTemplate jdbcTemplate;
public RLSConnectionManager(DataSource dataSource, JdbcTemplate jdbcTemplate) { public RLSConnectionManager(DataSource dataSource) {
this.dataSource = dataSource; this.dataSource = dataSource;
this.jdbcTemplate = jdbcTemplate;
} }
/** /**
@@ -43,18 +43,27 @@ public class RLSConnectionManager {
// CRITICAL: Set connection to manual commit to ensure atomicity // CRITICAL: Set connection to manual commit to ensure atomicity
connection.setAutoCommit(false); 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 { try {
// Set the RLS context variable on THIS connection // Set the RLS context variable on THIS connection using raw Statement
scopedTemplate.execute("SET LOCAL app.current_user_id = " + userId); 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 // Execute the operation with the scoped template
T result = operation.apply(scopedTemplate); T result = operation.apply(scopedTemplate);
@@ -70,11 +79,10 @@ public class RLSConnectionManager {
throw new RuntimeException("Error executing RLS operation", e); throw new RuntimeException("Error executing RLS operation", e);
} finally { } finally {
// CRITICAL: Reset the context variable before returning connection to pool // 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 (Statement stmt = connection.createStatement()) {
try { stmt.execute("RESET app.current_user_id");
scopedTemplate.execute("RESET app.current_user_id"); System.out.println("RLS context reset");
} catch (Exception e) { } 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()); System.err.println("Warning: Failed to reset RLS context: " + e.getMessage());
} }
} }

View File

@@ -69,17 +69,20 @@ public class RLSTestController {
"SELECT current_setting('app.current_user_id', true)", "SELECT current_setting('app.current_user_id', true)",
String.class String.class
); );
result.put("user1Context", ctx); result.put("user1Context", ctx.isEmpty() ? "EMPTY" : ctx);
return null; return null;
}); });
// Check if context leaked (should be null or empty) // Check if context leaked (should be null or empty)
String leakedContext = null; String leakedContext;
try { try {
leakedContext = jdbcTemplate.queryForObject( leakedContext = jdbcTemplate.queryForObject(
"SELECT current_setting('app.current_user_id', true)", "SELECT current_setting('app.current_user_id', true)",
String.class String.class
); );
if (leakedContext == null || leakedContext.isEmpty()) {
leakedContext = "NOT_SET";
}
} catch (Exception e) { } catch (Exception e) {
// Expected - variable should not be set // Expected - variable should not be set
leakedContext = "NOT_SET"; leakedContext = "NOT_SET";
@@ -92,7 +95,7 @@ public class RLSTestController {
"SELECT current_setting('app.current_user_id', true)", "SELECT current_setting('app.current_user_id', true)",
String.class String.class
); );
result.put("user2Context", ctx); result.put("user2Context", ctx.isEmpty() ? "EMPTY" : ctx);
return null; return null;
}); });
@@ -102,6 +105,9 @@ public class RLSTestController {
"SELECT current_setting('app.current_user_id', true)", "SELECT current_setting('app.current_user_id', true)",
String.class String.class
); );
if (leakedContext == null || leakedContext.isEmpty()) {
leakedContext = "NOT_SET";
}
} catch (Exception e) { } catch (Exception e) {
leakedContext = "NOT_SET"; leakedContext = "NOT_SET";
} }
@@ -157,27 +163,43 @@ public class RLSTestController {
// Enable RLS // Enable RLS
jdbcTemplate.execute("ALTER TABLE documents ENABLE ROW LEVEL SECURITY"); 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 // 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(""" jdbcTemplate.execute("""
CREATE POLICY user_documents_policy ON documents CREATE POLICY user_documents_policy ON documents
FOR ALL 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 // Insert test data WITH RLS context set
jdbcTemplate.update( // 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 (?, ?, ?)", "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"User 1 Document", "Private content for user 1", 1L "User 1 Document", "Private content for user 1", 1L
); );
jdbcTemplate.update( scopedTemplate.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 (?, ?, ?)", "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"Another User 1 Doc", "More private content for user 1", 1L "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";
} }
} }

View File

@@ -2,6 +2,6 @@ spring.application.name=demo
# PostgreSQL Database Configuration # PostgreSQL Database Configuration
spring.datasource.url=jdbc:postgresql://localhost:5432/rls_test spring.datasource.url=jdbc:postgresql://localhost:5432/rls_test
spring.datasource.username=postgres spring.datasource.username=rls_test
spring.datasource.password=postgres spring.datasource.password=rls_test
spring.datasource.hikari.maximum-pool-size=10 spring.datasource.hikari.maximum-pool-size=10