RLS works
This commit is contained in:
@@ -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());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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";
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user