spring-rl-test/src/main/java/com/example/demo/RLSTestController.java

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.isEmpty() ? "EMPTY" : ctx);
            return null;
        });
        
        // Check if context leaked (should be null or empty)
        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";
        }
        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.isEmpty() ? "EMPTY" : ctx);
            return null;
        });
        
        // Check again
        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) {
            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");
        
        // 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 = 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 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
            );
            scopedTemplate.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(
                "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";
    }
}