omarmunoz/spring-rl-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

before fixing latency

Browse Source
This commit is contained in:
Omar Muñoz
2026-01-21 10:44:12 -06:00
parent 7e09621a7d
commit 5dc7f8b2dd
5 changed files with 181 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/main/java/com/example/demo/DemoApplication.java
View File

@@ -2,8 +2,10 @@ package com.example.demo;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.transaction.annotation.EnableTransactionManagement;
@SpringBootApplication
@EnableTransactionManagement
public class DemoApplication {
public static void main(String[] args) {

131
src/main/java/com/example/demo/DocumentController.java Normal file
View File

@@ -0,0 +1,131 @@
package com.example.demo;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import java.util.HashMap;
import java.util.Map;
@RestController
@RequestMapping("/api/documents")
public class DocumentController {
private final RLSConnectionManager rlsManager;
private final DocumentRepository documentRepository;
public DocumentController(RLSConnectionManager rlsManager, DocumentRepository documentRepository) {
this.rlsManager = rlsManager;
this.documentRepository = documentRepository;
}
/**
* CREATE - Insert a new document using repository
*/
@PostMapping
public ResponseEntity<Document> createDocument(HttpServletRequest request,
@RequestBody Document document) {
Long userId = (Long) request.getAttribute("userId");
// Force the user_id to match JWT user (prevent privilege escalation)
document.setUserId(userId);
return rlsManager.executeWithRLSContext(userId, () -> {
Document saved = documentRepository.save(document);
return ResponseEntity.status(HttpStatus.CREATED).body(saved);
});
}
/**
* READ - Get all documents (filtered by RLS to current user)
*/
@GetMapping
public ResponseEntity<Iterable<Document>> getAllDocuments(HttpServletRequest request) {
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, () -> {
Iterable<Document> documents = documentRepository.findAll();
return ResponseEntity.ok(documents);
});
}
/**
* READ - Get single document by ID (RLS ensures user owns it)
*/
@GetMapping("/{id}")
public ResponseEntity<Map<String, Object>> getDocumentById(HttpServletRequest request,
@PathVariable Long id) {
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, () -> {
return documentRepository.findById(id)
.<ResponseEntity<Map<String, Object>>>map(doc -> {
Map<String, Object> response = new HashMap<>();
response.put("document", doc);
return ResponseEntity.ok(response);
})
.orElseGet(() -> {
Map<String, Object> error = new HashMap<>();
error.put("error", "Document not found or access denied");
return ResponseEntity.status(HttpStatus.NOT_FOUND).body(error);
});
});
}
/**
* UPDATE - Update existing document (RLS ensures user owns it)
*/
@PutMapping("/{id}")
public ResponseEntity<Map<String, Object>> updateDocument(HttpServletRequest request,
@PathVariable Long id,
@RequestBody Document updatedDocument) {
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, () -> {
return documentRepository.findById(id)
.<ResponseEntity<Map<String, Object>>>map(existingDoc -> {
// Update fields
existingDoc.setTitle(updatedDocument.getTitle());
existingDoc.setContent(updatedDocument.getContent());
// Don't allow changing user_id
Document saved = documentRepository.save(existingDoc);
Map<String, Object> response = new HashMap<>();
response.put("document", saved);
return ResponseEntity.ok(response);
})
.orElseGet(() -> {
Map<String, Object> error = new HashMap<>();
error.put("error", "Document not found or access denied");
return ResponseEntity.status(HttpStatus.NOT_FOUND).body(error);
});
});
}
/**
* DELETE - Delete document (RLS ensures user owns it)
*/
@DeleteMapping("/{id}")
public ResponseEntity<Map<String, Object>> deleteDocument(HttpServletRequest request,
@PathVariable Long id) {
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, () -> {
return documentRepository.findById(id)
.<ResponseEntity<Map<String, Object>>>map(doc -> {
documentRepository.deleteById(id);
Map<String, Object> response = new HashMap<>();
response.put("success", true);
response.put("message", "Document deleted");
response.put("id", id);
return ResponseEntity.ok(response);
})
.orElseGet(() -> {
Map<String, Object> error = new HashMap<>();
error.put("error", "Document not found or access denied");
return ResponseEntity.status(HttpStatus.NOT_FOUND).body(error);
});
});
}
}

1
src/main/java/com/example/demo/FilterConfig.java
View File

@@ -22,6 +22,7 @@ public class FilterConfig {
registrationBean.addUrlPatterns("/api/rls-test/setup");
registrationBean.addUrlPatterns("/api/rls-test/documents/*");
registrationBean.addUrlPatterns("/api/rls-test/context/*");
registrationBean.addUrlPatterns("/api/documents/*");
registrationBean.setOrder(1);

98
src/main/java/com/example/demo/RLSConnectionManager.java
View File

@@ -1,95 +1,61 @@
package com.example.demo;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.SingleConnectionDataSource;
import org.springframework.stereotype.Component;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.function.Function;
import java.util.function.Supplier;
/**
* Manages PostgreSQL RLS context variables with proper connection scoping.
*
* CRITICAL: This ensures context variables are set and reset on the SAME connection
* to prevent leaking between concurrent requests when connections are pooled.
* Uses programmatic transaction management for better performance.
*/
@Component
public class RLSConnectionManager {
private final DataSource dataSource;
private final JdbcTemplate jdbcTemplate;
private final PlatformTransactionManager transactionManager;
public RLSConnectionManager(DataSource dataSource) {
this.dataSource = dataSource;
public RLSConnectionManager(JdbcTemplate jdbcTemplate, PlatformTransactionManager transactionManager) {
this.jdbcTemplate = jdbcTemplate;
this.transactionManager = transactionManager;
}
/**
* Executes an operation with RLS context variables set on a dedicated connection.
* The connection is obtained, configured, used, reset, and returned to pool - all atomically.
* Executes an operation with RLS context variables set.
* Uses programmatic transactions for better performance.
*
* @param userId The user ID to set in the context
* @param operation The operation to execute (receives JdbcTemplate bound to the connection)
* @param operation The operation to execute
* @param <T> Return type
* @return Result of the operation
*/
public <T> T executeWithRLSContext(Long userId, Function<JdbcTemplate, T> operation) {
Connection connection = null;
public <T> T executeWithRLSContext(Long userId, Supplier<T> operation) {
// Start transaction
DefaultTransactionDefinition def = new DefaultTransactionDefinition();
def.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRED);
TransactionStatus status = transactionManager.getTransaction(def);
try {
// Get a connection from the pool
connection = dataSource.getConnection();
// Set the RLS context variable (uses transaction connection)
jdbcTemplate.execute("SET LOCAL app.current_user_id = '" + userId + "'");
// CRITICAL: Set connection to manual commit to ensure atomicity
connection.setAutoCommit(false);
// Execute the operation
T result = operation.get();
try {
// Set the RLS context variable on THIS connection
try (Statement stmt = connection.createStatement()) {
stmt.execute("SET LOCAL app.current_user_id = '" + userId + "'");
}
// Create a JdbcTemplate bound to THIS specific connection
SingleConnectionDataSource singleConnectionDataSource =
new SingleConnectionDataSource(connection, true);
JdbcTemplate scopedTemplate = new JdbcTemplate(singleConnectionDataSource);
// Execute the operation with the scoped template
T result = operation.apply(scopedTemplate);
// Commit the transaction
connection.commit();
return result;
} catch (Exception e) {
// Rollback on error
connection.rollback();
throw new RuntimeException("Error executing RLS operation", e);
} finally {
// CRITICAL: Reset the context variable before returning connection to pool
try (Statement stmt = connection.createStatement()) {
stmt.execute("RESET app.current_user_id");
} catch (Exception e) {
// Log but don't throw - connection will still be returned to pool
System.err.println("Warning: Failed to reset RLS context: " + e.getMessage());
}
}
// Commit transaction (SET LOCAL auto-reverts after commit)
transactionManager.commit(status);
} catch (SQLException e) {
throw new RuntimeException("Failed to obtain database connection", e);
} finally {
// Return connection to pool
if (connection != null) {
try {
connection.setAutoCommit(true); // Restore default behavior
connection.close(); // Returns to pool
} catch (SQLException e) {
System.err.println("Error closing connection: " + e.getMessage());
}
}
return result;
} catch (Exception e) {
// Rollback on error
transactionManager.rollback(status);
throw new RuntimeException("Error executing RLS operation", e);
}
}
}

30
src/main/java/com/example/demo/RLSTestController.java
View File

@@ -29,10 +29,10 @@ public class RLSTestController {
// Get user ID from JWT (set by JwtFilter)
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
return rlsManager.executeWithRLSContext(userId, () -> {
// This query will only return documents the user has access to (via RLS policy)
String sql = "SELECT id, title, content, user_id FROM documents";
return scopedTemplate.queryForList(sql);
return jdbcTemplate.queryForList(sql);
});
}
@@ -43,9 +43,9 @@ public class RLSTestController {
public Map<String, Object> verifyContextVariable(HttpServletRequest request) {
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
return rlsManager.executeWithRLSContext(userId, () -> {
// Query the context variable to verify it's set
String currentUserId = scopedTemplate.queryForObject(
String currentUserId = jdbcTemplate.queryForObject(
"SELECT current_setting('app.current_user_id', true)",
String.class
);
@@ -68,8 +68,8 @@ public class RLSTestController {
Map<String, Object> result = new HashMap<>();
// Set context for user 1
rlsManager.executeWithRLSContext(1L, scopedTemplate -> {
String ctx = scopedTemplate.queryForObject(
rlsManager.executeWithRLSContext(1L, () -> {
String ctx = jdbcTemplate.queryForObject(
"SELECT current_setting('app.current_user_id', true)",
String.class
);
@@ -94,8 +94,8 @@ public class RLSTestController {
result.put("afterUser1", leakedContext);
// Set context for user 2
rlsManager.executeWithRLSContext(2L, scopedTemplate -> {
String ctx = scopedTemplate.queryForObject(
rlsManager.executeWithRLSContext(2L, () -> {
String ctx = jdbcTemplate.queryForObject(
"SELECT current_setting('app.current_user_id', true)",
String.class
);
@@ -133,9 +133,9 @@ public class RLSTestController {
// Get user ID from JWT
Long userId = (Long) request.getAttribute("userId");
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
return rlsManager.executeWithRLSContext(userId, () -> {
// Insert with the user context set
scopedTemplate.update(
jdbcTemplate.update(
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
title, content, userId
);
@@ -188,20 +188,20 @@ public class RLSTestController {
// 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(
rlsManager.executeWithRLSContext(1L, () -> {
jdbcTemplate.update(
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"User 1 Document", "Private content for user 1", 1L
);
scopedTemplate.update(
jdbcTemplate.update(
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"Another User 1 Doc", "More private content for user 1", 1L
);
return null;
});
rlsManager.executeWithRLSContext(2L, scopedTemplate -> {
scopedTemplate.update(
rlsManager.executeWithRLSContext(2L, () -> {
jdbcTemplate.update(
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
"User 2 Document", "Private content for user 2", 2L
);