RLS initial test
This commit is contained in:
Omar Muñoz
183
src/main/java/com/example/demo/RLSTestController.java
Normal file
183
src/main/java/com/example/demo/RLSTestController.java
Normal file
@@ -0,0 +1,183 @@
|
||||
package com.example.demo;
|
||||
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.springframework.web.bind.annotation.*;
|
||||
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
@RestController
|
||||
@RequestMapping("/api/rls-test")
|
||||
public class RLSTestController {
|
||||
|
||||
private final RLSConnectionManager rlsManager;
|
||||
private final DocumentRepository documentRepository;
|
||||
private final JdbcTemplate jdbcTemplate;
|
||||
|
||||
public RLSTestController(RLSConnectionManager rlsManager,
|
||||
DocumentRepository documentRepository,
|
||||
JdbcTemplate jdbcTemplate) {
|
||||
this.rlsManager = rlsManager;
|
||||
this.documentRepository = documentRepository;
|
||||
this.jdbcTemplate = jdbcTemplate;
|
||||
}
|
||||
|
||||
/**
|
||||
* Test 1: Execute raw SQL with RLS context
|
||||
*/
|
||||
@GetMapping("/documents/user/{userId}")
|
||||
public List<Map<String, Object>> getDocumentsWithRawSQL(@PathVariable Long userId) {
|
||||
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
|
||||
// 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);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Test 2: Verify context variable is set correctly
|
||||
*/
|
||||
@GetMapping("/context/verify/{userId}")
|
||||
public Map<String, Object> verifyContextVariable(@PathVariable Long userId) {
|
||||
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
|
||||
// Query the context variable to verify it's set
|
||||
String currentUserId = scopedTemplate.queryForObject(
|
||||
"SELECT current_setting('app.current_user_id', true)",
|
||||
String.class
|
||||
);
|
||||
|
||||
Map<String, Object> result = new HashMap<>();
|
||||
result.put("requestedUserId", userId);
|
||||
result.put("contextUserId", currentUserId);
|
||||
result.put("match", userId.toString().equals(currentUserId));
|
||||
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Test 3: Verify context is reset after request (simulate concurrent requests)
|
||||
*/
|
||||
@GetMapping("/context/isolation-test")
|
||||
public Map<String, Object> testContextIsolation() throws InterruptedException {
|
||||
Map<String, Object> result = new HashMap<>();
|
||||
|
||||
// Set context for user 1
|
||||
rlsManager.executeWithRLSContext(1L, scopedTemplate -> {
|
||||
String ctx = scopedTemplate.queryForObject(
|
||||
"SELECT current_setting('app.current_user_id', true)",
|
||||
String.class
|
||||
);
|
||||
result.put("user1Context", ctx);
|
||||
return null;
|
||||
});
|
||||
|
||||
// Check if context leaked (should be null or empty)
|
||||
String leakedContext = null;
|
||||
try {
|
||||
leakedContext = jdbcTemplate.queryForObject(
|
||||
"SELECT current_setting('app.current_user_id', true)",
|
||||
String.class
|
||||
);
|
||||
} catch (Exception e) {
|
||||
// Expected - variable should not be set
|
||||
leakedContext = "NOT_SET";
|
||||
}
|
||||
result.put("afterUser1", leakedContext);
|
||||
|
||||
// Set context for user 2
|
||||
rlsManager.executeWithRLSContext(2L, scopedTemplate -> {
|
||||
String ctx = scopedTemplate.queryForObject(
|
||||
"SELECT current_setting('app.current_user_id', true)",
|
||||
String.class
|
||||
);
|
||||
result.put("user2Context", ctx);
|
||||
return null;
|
||||
});
|
||||
|
||||
// Check again
|
||||
try {
|
||||
leakedContext = jdbcTemplate.queryForObject(
|
||||
"SELECT current_setting('app.current_user_id', true)",
|
||||
String.class
|
||||
);
|
||||
} catch (Exception e) {
|
||||
leakedContext = "NOT_SET";
|
||||
}
|
||||
result.put("afterUser2", leakedContext);
|
||||
|
||||
result.put("isolationSuccess", "NOT_SET".equals(result.get("afterUser1")) &&
|
||||
"NOT_SET".equals(result.get("afterUser2")));
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Test 4: Insert with RLS context (useful for audit trails)
|
||||
*/
|
||||
@PostMapping("/documents")
|
||||
public Map<String, Object> createDocument(@RequestParam Long userId,
|
||||
@RequestParam String title,
|
||||
@RequestParam String content) {
|
||||
return rlsManager.executeWithRLSContext(userId, scopedTemplate -> {
|
||||
// Insert with the user context set
|
||||
scopedTemplate.update(
|
||||
"INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
|
||||
title, content, userId
|
||||
);
|
||||
|
||||
Map<String, Object> result = new HashMap<>();
|
||||
result.put("success", true);
|
||||
result.put("userId", userId);
|
||||
result.put("title", title);
|
||||
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Setup endpoint - creates the table and test data
|
||||
*/
|
||||
@PostMapping("/setup")
|
||||
public String setupDatabase() {
|
||||
// Drop existing table
|
||||
jdbcTemplate.execute("DROP TABLE IF EXISTS documents CASCADE");
|
||||
|
||||
// Create table
|
||||
jdbcTemplate.execute("""
|
||||
CREATE TABLE documents (
|
||||
id SERIAL PRIMARY KEY,
|
||||
title VARCHAR(255) NOT NULL,
|
||||
content TEXT,
|
||||
user_id BIGINT NOT NULL
|
||||
)
|
||||
""");
|
||||
|
||||
// Enable RLS
|
||||
jdbcTemplate.execute("ALTER TABLE documents ENABLE ROW LEVEL SECURITY");
|
||||
|
||||
// Create RLS policy
|
||||
jdbcTemplate.execute("""
|
||||
CREATE POLICY user_documents_policy ON documents
|
||||
FOR ALL
|
||||
USING (user_id = 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
|
||||
);
|
||||
|
||||
return "Database setup complete with RLS enabled";
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user