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;
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());
}
}

View File

@@ -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 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
);
jdbcTemplate.update(
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"User 2 Document", "Private content for user 2", 2L
);
jdbcTemplate.update(
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";
}
}

View File

@@ -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