diff --git a/CLAUDE.md b/CLAUDE.md index 653c3ff..b046d0e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -199,8 +199,9 @@ returns a number). 2. **LLM Provider Registry**: Use `LlmProviderRegistry` for all provider-specific LLM behavior. Do NOT add if/else or switch on provider type. Chat and embedding providers are registered and resolved independently — some providers offer only one. Providers are *factories* over credentials, not `ChatModel`s, so credentials stay resolvable per call and key rotation needs no restart. 3. **SSH-Aware Access**: Always use `ConnectionService.getJdbcTemplate(connectionId, request)` — handles SSH tunneling transparently. 4. **SQL Rule**: All generated SQL MUST use table-qualified column names (`table.column_name`). -5. **RAG Caching**: Three-tier cache (memory → Redis → Azure Search). Redis failure is graceful (app continues without caching). -6. **Virtual Threads**: Enabled for concurrency (JDK 25). +5. **Chat access policy**: Fail closed. Walk the whole SQL tree (CTEs, set ops, subqueries). Deny unparseable or unhandled statements. Require an actor except `INTERNAL`/`SCHEDULED`. MCP/Editor identity comes from `SecurityContext`, not `QueryActorContextHolder`. Persist `allowed_schemas`. Do not let "how many" override a protected-column mention. Public share is refused when the connection has an active policy. +6. **RAG Caching**: Three-tier cache (memory → Redis → Azure Search). Redis failure is graceful (app continues without caching). +7. **Virtual Threads**: Enabled for concurrency (JDK 25). ### Frontend Rules 1. **API Centralization**: ALL API calls through `src/lib/api/client.js`. Never create direct axios instances. @@ -210,7 +211,13 @@ returns a number). 5. **Design**: Minimal black/white/grey palette, Inter font, subtle transitions. See UX guidelines in full CLAUDE.md. ### Admin profile switch -Admins can **View as** a sub-user from the top-right of the home layout (`ProfileSwitch`) to verify connection ACLs, chat/editor policies, and role-gated nav. The admin JWT stays on the session; `ImpersonationService` sets an httpOnly `impersonate_user` cookie and `JwtAuthenticationFilter` overlays the target principal. `POST|DELETE|GET /api/admin/impersonate` are excluded from the overlay so stop/list still run as the real admin. Cannot target another ADMIN, self, or a non-ACTIVE account. `/auth/me` returns the **effective** user plus `impersonating` / `impersonatorUsername`. +Admins can **View as** a sub-user from the top-right of the home layout (`ProfileSwitch`) to verify connection ACLs, chat/editor policies, and role-gated nav. + +The admin JWT **subject** stays the administrator so logout, refresh, and `/admin/impersonate` still own the real session. Policy identity is the target: an httpOnly `impersonate_user` cookie plus an `impUid` claim on the access token. `JwtAuthenticationFilter` overlays that principal onto the SecurityContext for every request except the impersonation control plane, logout, and session refresh. Chat, Editor, schema listing, and Agent MCP calls then run `AccessControlService` / `ConnectionChatAccessPolicyService` as the target (`actorIsAdmin` is false, so policies apply). + +The Agent tab must not inherit the admin MCP token. `/api/agent/session` mints an MCP token for the effective user and never falls back to the admin session JWT while View as is active. nginx `auth_request` on `/agent-api` forwards `/api/auth/me`'s `X-Remote-User` (the overlaid username) instead of hardcoding `admin`. + +`POST|DELETE|GET /api/admin/impersonate` are excluded from the overlay so stop/list still run as the real admin. Cannot target another ADMIN, self, or a non-ACTIVE account. `/auth/me` returns the **effective** user plus `impersonating` / `impersonatorUsername`. ### Git Rules - Do NOT commit automatically — wait for explicit user instruction. @@ -238,7 +245,16 @@ Admins can **View as** a sub-user from the top-right of the home layout (`Profil `hermes_requires <0.20.0` on a "verified" 401 that came from a hand-rolled `hermes serve` run rather than `hermes webui`. 0.20.0 works. Verify against the real start path before writing a version constraint. -5. **The agent image build clones two third-party repos over the public internet, +5. **Hermes MCP is process-global.** One `deepsql-phase1-server.js` stdio server + is started from the first loaded profile (usually `u-admin`) and lives until + the Agent API process dies. `POST /api/profile/switch` is `process_wide=False` + on purpose (per-browser cookie), so a View as / new-thread Agent chat tagged + `profile: u-marts-editor` still sends the admin MCP bearer. Chat-access policy + then `resolveEffectivePolicy(..., actorIsAdmin=true)` → `none()`. The + provisioner mirrors the target user's token onto every `deepsql.token` the + live process might re-read (`DEEPSQL_TOKEN_FILE` mtime cache in + `mcp/deepsql-phase1-lib.js`). +6. **The agent image build clones two third-party repos over the public internet, unauthenticated.** `agent/Dockerfile` fetches `NousResearch/hermes-agent` and `nesquena/hermes-webui` at build time. GitHub rate-limits unauthenticated requests *per source IP*, and Actions runners share pooled egress addresses, so @@ -276,6 +292,16 @@ broken. Assert the *outcome*, never the attempt: - **`set -e` + `read` at EOF aborts silently.** Prompts in `install.sh` use `read … || true` so the explicit emptiness checks report the problem. Without it the installer exited 1 with no message, after writing generated secrets to `.env`. +- **Minting an Agent MCP token ≠ Hermes using it.** `/api/agent/session` can mint + `u-marts-editor`'s token and `POST /api/profile/switch` can 200 while + `execute_sql` still authenticates as admin. Hermes keeps one DeepSQL MCP + stdio process (started from the first profile that loaded `mcp_servers`) and + `profile/switch` is `process_wide=False`. A new chat thread does not respawn + MCP. `probeMcpAuth` only proves the *minted* token works against Spring, not + that the live MCP process will send it. The provisioner must mirror the + token onto every `deepsql.token` the live process might be watching + (`scripts/local-agent-provisioner.py`). Audit: `security_event.user_id` on + `EDITOR_QUERY_EXECUTED` with `clientType=mcp`. - **Silent-failure rule, concretely:** the CLI rendered an unreachable server as `No databases connected yet` because one `catch` covered both the connection fetch and decorative extras. An unreachable host must never look like an empty account. diff --git a/backend/src/main/java/com/dbaagent/controller/AuthController.java b/backend/src/main/java/com/dbaagent/controller/AuthController.java index e5d3cf7..105b272 100644 --- a/backend/src/main/java/com/dbaagent/controller/AuthController.java +++ b/backend/src/main/java/com/dbaagent/controller/AuthController.java @@ -207,11 +207,21 @@ public ResponseEntity refreshSession(HttpServletRequest httpRequest, HttpServ User effectiveUser = impersonationService.resolveFromCookie(httpRequest, user) .map(ImpersonationContext.State::target) .orElse(user); - // Keep the user's agent token alive for as long as the UI session lives. - // The SPA refreshes on access-token expiry (~every 15 min of activity), so - // this slides the agent token forward on each active interval — a logged-in - // UI never ends up with a dead agent. + if (effectiveUser != user && effectiveUser.getId() != null) { + authSessionService.reissueAccessToken( + httpResponse, + session.getId(), + user, + effectiveUser.getId() + ); + } + // Keep agent tokens alive for as long as the UI session lives. + // During View as the SPA still refreshes the *admin* session; also + // slide the target user's minted MCP token or their Agent tab dies. agentBridgeService.extendAgentTokens(user.getUsername()); + if (!effectiveUser.getUsername().equals(user.getUsername())) { + agentBridgeService.extendAgentTokens(effectiveUser.getUsername()); + } Map payload = toAuthPayload( effectiveUser, effectiveUser.getRoleEnum(), diff --git a/backend/src/main/java/com/dbaagent/controller/BrainController.java b/backend/src/main/java/com/dbaagent/controller/BrainController.java index 636cdc3..6aabfdb 100644 --- a/backend/src/main/java/com/dbaagent/controller/BrainController.java +++ b/backend/src/main/java/com/dbaagent/controller/BrainController.java @@ -65,6 +65,7 @@ import com.dbaagent.service.brain.query.PlanPatternLibraryService; import com.dbaagent.service.brain.BrainInsightEmbeddingService; import com.dbaagent.service.SlowQueryHistoryService; +import com.dbaagent.service.UserDataAccessPolicyService; import com.dbaagent.service.security.AccessControlService; import com.dbaagent.model.SlowQueryAnalysis; import lombok.RequiredArgsConstructor; @@ -139,6 +140,7 @@ public class BrainController { private final SlowQueryHistoryService slowQueryHistoryService; private final BrainInsightEmbeddingService brainInsightEmbeddingService; private final AccessControlService accessControlService; + private final UserDataAccessPolicyService userDataAccessPolicyService; @GetMapping("/understanding/{connectionId}") public ResponseEntity getUnderstanding( @@ -531,7 +533,12 @@ public ResponseEntity> getInferredRelationships( ) { try { accessControlService.assertCanReadConnectionContent(connectionId); - return ResponseEntity.ok(joinRelationshipInferenceService.getRelationships(connectionId)); + return ResponseEntity.ok(userDataAccessPolicyService.filterInferredRelationships( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + joinRelationshipInferenceService.getRelationships(connectionId) + )); } catch (ResponseStatusException e) { throw e; } catch (Exception e) { diff --git a/backend/src/main/java/com/dbaagent/controller/DashboardQueryController.java b/backend/src/main/java/com/dbaagent/controller/DashboardQueryController.java index 922039f..7022ec9 100644 --- a/backend/src/main/java/com/dbaagent/controller/DashboardQueryController.java +++ b/backend/src/main/java/com/dbaagent/controller/DashboardQueryController.java @@ -77,7 +77,10 @@ public ResponseEntity query(@RequestBody DashboardQueryRequest request) { qr.setExecutionOrigin(QueryExecutionOrigin.API); QueryResult result = queryExecutorService.executeQuery( request.connectionId(), qr, - QueryExecutionContext.api(accessControlService.getCurrentUsername())); + QueryExecutionContext.api( + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin() + )); return ResponseEntity.ok(Map.of( "success", true, "columns", result.getColumns() == null ? java.util.List.of() : result.getColumns(), diff --git a/backend/src/main/java/com/dbaagent/controller/McpController.java b/backend/src/main/java/com/dbaagent/controller/McpController.java index 793bf1c..243b0cf 100644 --- a/backend/src/main/java/com/dbaagent/controller/McpController.java +++ b/backend/src/main/java/com/dbaagent/controller/McpController.java @@ -8,7 +8,6 @@ import com.dbaagent.model.QueryResult; import com.dbaagent.service.ExplainPlanService; import com.dbaagent.service.McpSqlGuardService; -import com.dbaagent.service.QueryActorContextHolder; import com.dbaagent.service.QueryExecutionContext; import com.dbaagent.service.QueryExecutionPolicyException; import com.dbaagent.service.QueryExecutorService; @@ -77,7 +76,10 @@ public ResponseEntity executeReadOnlyQuery(@RequestBody McpReadOnlyQueryReque QueryResult result = queryExecutorService.executeQuery( request.getConnectionId(), queryRequest, - QueryExecutionContext.mcp(QueryActorContextHolder.currentUsername()) + QueryExecutionContext.mcp( + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin() + ) ); return ResponseEntity.ok(Map.of( "success", true, diff --git a/backend/src/main/java/com/dbaagent/controller/SavedDashboardController.java b/backend/src/main/java/com/dbaagent/controller/SavedDashboardController.java index c2e3364..7dc8ca5 100644 --- a/backend/src/main/java/com/dbaagent/controller/SavedDashboardController.java +++ b/backend/src/main/java/com/dbaagent/controller/SavedDashboardController.java @@ -2,6 +2,7 @@ import com.dbaagent.model.DashboardVersion; import com.dbaagent.model.SavedDashboard; +import com.dbaagent.service.ConnectionChatAccessPolicyService; import com.dbaagent.service.SavedDashboardService; import com.dbaagent.service.security.AccessControlService; import lombok.extern.slf4j.Slf4j; @@ -27,6 +28,9 @@ public class SavedDashboardController { @Autowired private AccessControlService accessControlService; + @Autowired + private ConnectionChatAccessPolicyService connectionChatAccessPolicyService; + // Every write method below is load-then-save on a row a background generation // turn (SavedDashboardService.beginGenerationTurn etc.) may be writing at the // same time. Without this helper, the loser's raw Hibernate message @@ -47,6 +51,13 @@ public ResponseEntity> enableShare(@PathVariable UUID id) { SavedDashboard existing = savedDashboardService.getDashboardById(id) .orElseThrow(() -> new IllegalArgumentException("Dashboard not found")); accessControlService.assertCanReadConnectionContent(existing.getConnectionId()); + if (connectionChatAccessPolicyService.hasActivePolicy(existing.getConnectionId())) { + return ResponseEntity.status(HttpStatus.CONFLICT).body(Map.of( + "success", false, + "errorCode", "POLICY_PUBLIC_SHARE_FORBIDDEN", + "message", "This connection has an active chat access policy, so the dashboard cannot be shared publicly." + )); + } SavedDashboard d = savedDashboardService.enablePublicShare(id); return ResponseEntity.ok(Map.of("success", true, "shareToken", d.getShareToken(), "isPublic", true)); diff --git a/backend/src/main/java/com/dbaagent/controller/SchemaController.java b/backend/src/main/java/com/dbaagent/controller/SchemaController.java index ca3064f..6f37ab9 100644 --- a/backend/src/main/java/com/dbaagent/controller/SchemaController.java +++ b/backend/src/main/java/com/dbaagent/controller/SchemaController.java @@ -7,6 +7,7 @@ import com.dbaagent.service.QueryExecutionContext; import com.dbaagent.service.QueryExecutionPolicyException; import com.dbaagent.service.ActiveQueryService; +import com.dbaagent.service.McpTokenService; import com.dbaagent.service.QueryExecutorService; import com.dbaagent.service.RunningQueryRegistry; import com.dbaagent.service.SqlExecutionAuditService; @@ -15,6 +16,7 @@ import com.dbaagent.service.SchemaScannerService; import com.dbaagent.service.VisualizationService; import com.dbaagent.service.security.AccessControlService; +import org.springframework.http.HttpHeaders; import jakarta.servlet.http.HttpServletRequest; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -195,11 +197,7 @@ public ResponseEntity> executeQuery( QueryResult result = queryExecutorService.executeQuery( connectionId, queryRequest, - QueryExecutionContext.editor( - accessControlService.getCurrentUsername(), - accessControlService.isCurrentUserAdmin(), - Boolean.TRUE.equals(queryRequest.getMutationConfirmed()) - ) + queryExecutionContext(queryRequest, httpRequest) ); sqlExecutionAuditService.record(SqlExecutionAuditService.AuditRecord.executed() .connectionId(connectionId) @@ -413,6 +411,26 @@ public ResponseEntity> getTableStats( } } + private QueryExecutionContext queryExecutionContext(QueryRequest queryRequest, HttpServletRequest httpRequest) { + String username = accessControlService.getCurrentUsername(); + boolean admin = accessControlService.isCurrentUserAdmin(); + if (isMcpBearer(httpRequest)) { + return QueryExecutionContext.mcp(username, admin); + } + return QueryExecutionContext.editor( + username, + admin, + Boolean.TRUE.equals(queryRequest.getMutationConfirmed()) + ); + } + + private boolean isMcpBearer(HttpServletRequest httpRequest) { + String authorization = httpRequest.getHeader(HttpHeaders.AUTHORIZATION); + return authorization != null + && authorization.startsWith("Bearer ") + && authorization.substring(7).startsWith(McpTokenService.TOKEN_PREFIX); + } + private SchemaMetadata scopedSchema(String connectionId, SchemaMetadata schema) { return userDataAccessPolicyService.filterSchemaMetadata( connectionId, diff --git a/backend/src/main/java/com/dbaagent/controller/TrainingController.java b/backend/src/main/java/com/dbaagent/controller/TrainingController.java index 59c3b53..4c5d7a6 100644 --- a/backend/src/main/java/com/dbaagent/controller/TrainingController.java +++ b/backend/src/main/java/com/dbaagent/controller/TrainingController.java @@ -10,6 +10,7 @@ import com.dbaagent.service.SemanticModelService; import com.dbaagent.service.TrainingJobService; import com.dbaagent.service.TrainingService; +import com.dbaagent.service.UserDataAccessPolicyService; import com.dbaagent.service.security.AccessControlService; import lombok.Data; import lombok.RequiredArgsConstructor; @@ -37,6 +38,7 @@ public class TrainingController { private final CredentialRepository credentialRepository; private final SemanticModelService semanticModelService; private final AccessControlService accessControlService; + private final UserDataAccessPolicyService userDataAccessPolicyService; /** * Train with schema DDL @@ -186,7 +188,9 @@ public ResponseEntity> debugRetrieve( if (question == null || question.isBlank()) { return ResponseEntity.badRequest().body(Map.of("error", "Query parameter 'q' is required")); } - return ResponseEntity.ok(trainingService.debugRetrieve(connectionId, question, topK)); + Map payload = trainingService.debugRetrieve(connectionId, question, topK); + filterDebugRetrieval(connectionId, payload); + return ResponseEntity.ok(payload); } catch (org.springframework.web.server.ResponseStatusException e) { throw e; } catch (Exception e) { @@ -346,6 +350,34 @@ private void rebuildAndReindexConnection(String connectionId) { trainingService.reindexConnection(connectionId); } + @SuppressWarnings("unchecked") + private void filterDebugRetrieval(String connectionId, Map payload) { + if (payload == null) { + return; + } + Object results = payload.get("results"); + if (!(results instanceof List rows)) { + return; + } + List> filtered = new java.util.ArrayList<>(); + for (Object row : rows) { + if (!(row instanceof Map map)) { + continue; + } + Object metadata = map.get("metadata"); + if (userDataAccessPolicyService.isRagMetadataInScope( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + metadata == null ? null : String.valueOf(metadata) + )) { + filtered.add((Map) map); + } + } + payload.put("results", filtered); + payload.put("resultCount", filtered.size()); + } + @Data public static class DocumentationRequest { private String connectionId; diff --git a/backend/src/main/java/com/dbaagent/dto/ConnectionChatAccessPolicyResponse.java b/backend/src/main/java/com/dbaagent/dto/ConnectionChatAccessPolicyResponse.java index 2da7d69..0373fcf 100644 --- a/backend/src/main/java/com/dbaagent/dto/ConnectionChatAccessPolicyResponse.java +++ b/backend/src/main/java/com/dbaagent/dto/ConnectionChatAccessPolicyResponse.java @@ -16,6 +16,8 @@ public class ConnectionChatAccessPolicyResponse { List blockedSensitivityCategories; List deniedTables; List deniedColumns; + List allowedSchemas; + boolean allowAggregates; boolean blockMode; boolean redactMode; boolean active; diff --git a/backend/src/main/java/com/dbaagent/dto/PolicyPreviewResponse.java b/backend/src/main/java/com/dbaagent/dto/PolicyPreviewResponse.java index 080ea28..6d71b16 100644 --- a/backend/src/main/java/com/dbaagent/dto/PolicyPreviewResponse.java +++ b/backend/src/main/java/com/dbaagent/dto/PolicyPreviewResponse.java @@ -11,6 +11,8 @@ public class PolicyPreviewResponse { List blockedSensitivityCategories; List deniedTables; List deniedColumns; + List allowedSchemas; + boolean allowAggregates; List impactedTables; List impactedColumns; boolean blockMode; diff --git a/backend/src/main/java/com/dbaagent/model/ConnectionChatAccessPolicy.java b/backend/src/main/java/com/dbaagent/model/ConnectionChatAccessPolicy.java index f1aa6e1..cc48f8c 100644 --- a/backend/src/main/java/com/dbaagent/model/ConnectionChatAccessPolicy.java +++ b/backend/src/main/java/com/dbaagent/model/ConnectionChatAccessPolicy.java @@ -40,6 +40,13 @@ public class ConnectionChatAccessPolicy { @Column(name = "denied_columns", columnDefinition = "jsonb") private List deniedColumns; + @JdbcTypeCode(SqlTypes.JSON) + @Column(name = "allowed_schemas", columnDefinition = "jsonb") + private List allowedSchemas; + + @Column(name = "allow_aggregates", nullable = false) + private boolean allowAggregates = false; + @Column(name = "block_mode", nullable = false) private boolean blockMode = true; diff --git a/backend/src/main/java/com/dbaagent/repository/ConnectionChatAccessPolicyRepository.java b/backend/src/main/java/com/dbaagent/repository/ConnectionChatAccessPolicyRepository.java index c6e4f35..76bf572 100644 --- a/backend/src/main/java/com/dbaagent/repository/ConnectionChatAccessPolicyRepository.java +++ b/backend/src/main/java/com/dbaagent/repository/ConnectionChatAccessPolicyRepository.java @@ -9,5 +9,6 @@ public interface ConnectionChatAccessPolicyRepository extends JpaRepository { Optional findByConnectionIdAndUsernameIgnoreCase(String connectionId, String username); List findAllByUsernameIgnoreCaseOrderByUpdatedAtDesc(String username); + boolean existsByConnectionIdAndActiveTrue(String connectionId); void deleteByConnectionId(String connectionId); } diff --git a/backend/src/main/java/com/dbaagent/security/ImpersonationContext.java b/backend/src/main/java/com/dbaagent/security/ImpersonationContext.java index 240354f..11d9cfe 100644 --- a/backend/src/main/java/com/dbaagent/security/ImpersonationContext.java +++ b/backend/src/main/java/com/dbaagent/security/ImpersonationContext.java @@ -5,11 +5,12 @@ import java.util.Optional; /** - * Request-scoped impersonation overlay. The admin JWT stays on the session; + * Request-scoped impersonation overlay. The admin JWT subject stays on the + * session so logout/refresh/control-plane still own the real administrator; * {@link JwtAuthenticationFilter} swaps the SecurityContext principal to the - * target user and records both identities here so {@code /auth/me} can show a - * banner and {@code AccessControlService} can honour the target even when - * {@code security.auth.enabled} is false. + * target user (from {@code impersonate_user} cookie or {@code impUid} claim) + * so policy evaluation, {@code /auth/me}, and {@code AccessControlService} + * honour the viewed-as user — including when {@code security.auth.enabled} is false. */ public final class ImpersonationContext { diff --git a/backend/src/main/java/com/dbaagent/security/JwtAuthenticationFilter.java b/backend/src/main/java/com/dbaagent/security/JwtAuthenticationFilter.java index 9262c33..0df39d6 100644 --- a/backend/src/main/java/com/dbaagent/security/JwtAuthenticationFilter.java +++ b/backend/src/main/java/com/dbaagent/security/JwtAuthenticationFilter.java @@ -143,12 +143,30 @@ private void applyImpersonation(HttpServletRequest request, HttpServletResponse throws ServletException, IOException { try { impersonationService.applyToRequest(request); + stampEffectiveUser(response); chain.doFilter(request, response); } finally { ImpersonationContext.clear(); } } + /** + * nginx {@code auth_request} on {@code /agent-api} forwards this as + * {@code X-Remote-User}. It must be the effective principal so + * View as (and non-admin Agent users) do not run the shared admin profile. + */ + private void stampEffectiveUser(HttpServletResponse response) { + var authentication = SecurityContextHolder.getContext().getAuthentication(); + if (authentication == null || !authentication.isAuthenticated()) { + return; + } + String name = authentication.getName(); + if (name == null || name.isBlank() || "anonymousUser".equals(name)) { + return; + } + response.setHeader("X-Remote-User", name); + } + private String extractUsernameSafely(String token) { if (token == null || token.isBlank()) { return null; diff --git a/backend/src/main/java/com/dbaagent/security/JwtUtil.java b/backend/src/main/java/com/dbaagent/security/JwtUtil.java index 9a3d206..1cdfc99 100644 --- a/backend/src/main/java/com/dbaagent/security/JwtUtil.java +++ b/backend/src/main/java/com/dbaagent/security/JwtUtil.java @@ -134,6 +134,17 @@ public String generateAccessToken( Role role, Set permissions, Duration ttl + ) { + return generateAccessToken(username, sessionId, role, permissions, ttl, null); + } + + public String generateAccessToken( + String username, + String sessionId, + Role role, + Set permissions, + Duration ttl, + Long impersonateUserId ) { Map claims = new HashMap<>(); claims.put("role", role.name()); @@ -145,10 +156,36 @@ public String generateAccessToken( if (sessionId != null && !sessionId.isBlank()) { claims.put("sid", sessionId); } + if (impersonateUserId != null && impersonateUserId > 0) { + claims.put("impUid", impersonateUserId); + } return createToken(claims, username, ttl); } + /** + * Target user id stamped onto an admin access token during View as. + * The JWT subject stays the administrator so logout/refresh/control-plane + * still own the real session; policy evaluation overlays this user. + */ + public Long extractImpersonateUserId(String token) { + Claims claims = extractAllClaims(token); + Object raw = claims.get("impUid"); + if (raw instanceof Number number) { + long value = number.longValue(); + return value > 0 ? value : null; + } + if (raw instanceof String text && !text.isBlank()) { + try { + long value = Long.parseLong(text.trim()); + return value > 0 ? value : null; + } catch (NumberFormatException e) { + return null; + } + } + return null; + } + /** * Generate token with role string and permissions set. */ diff --git a/backend/src/main/java/com/dbaagent/service/AgentBridgeService.java b/backend/src/main/java/com/dbaagent/service/AgentBridgeService.java index 5c75cbd..ec3d68c 100644 --- a/backend/src/main/java/com/dbaagent/service/AgentBridgeService.java +++ b/backend/src/main/java/com/dbaagent/service/AgentBridgeService.java @@ -1,5 +1,6 @@ package com.dbaagent.service; +import com.dbaagent.security.ImpersonationContext; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -142,9 +143,19 @@ public record ProfileBootstrap(String profile, String token) {} public ProfileBootstrap ensureProfile(String username, String authToken, String connectionId) { String profile = profileFor(username); if (!provisionEnabled) { + if (ImpersonationContext.isActive()) { + throw new ProvisioningException( + "Agent provisioning is disabled; View as cannot bind a user-scoped Agent token" + ); + } return new ProfileBootstrap(profile, authToken); } if (provisionSecret == null || provisionSecret.isBlank()) { + if (ImpersonationContext.isActive()) { + throw new ProvisioningException( + "Agent provisioning is not configured; View as cannot bind a user-scoped Agent token" + ); + } log.warn("agent.provision-secret is unset — skipping per-user provisioning for {}", username); return new ProfileBootstrap(profile, authToken); } @@ -154,11 +165,21 @@ public ProfileBootstrap ensureProfile(String username, String authToken, String // every MCP call. Mint a dedicated, user-scoped, revocable MCP token // instead (authenticated by McpTokenAuthenticationFilter, not the session // filter). Fall back to the session token only if minting fails so the - // tab still opens. + // tab still opens — except during View as, where the session JWT is the + // administrator's and would skip the target user's policy. String agentToken = mintAgentToken(username); if (agentToken == null) { + if (ImpersonationContext.isActive()) { + throw new ProvisioningException( + "Could not mint an MCP token for " + username + " while viewing as that user" + ); + } agentToken = authToken == null ? "" : authToken; } + // The provisioner also mirrors this token onto every profile + // `deepsql.token` the already-running Hermes MCP subprocess watches. + // Profile switch does not respawn MCP; without that mirror, View as + // Agent chats keep the admin credential and skip policy. callProvisioner(username, profile, agentToken, connectionId); return new ProfileBootstrap(profile, agentToken); } diff --git a/backend/src/main/java/com/dbaagent/service/AgentChatClient.java b/backend/src/main/java/com/dbaagent/service/AgentChatClient.java index 2f4a216..d7a5df3 100644 --- a/backend/src/main/java/com/dbaagent/service/AgentChatClient.java +++ b/backend/src/main/java/com/dbaagent/service/AgentChatClient.java @@ -136,7 +136,7 @@ private void switchProfile(String profile) throws Exception { throw new IllegalArgumentException("agent profile is required"); } // Profiles are provisioned as u-; the trusted-auth header is the - // bare username (nginx hard-codes X-Remote-User: admin for the browser path). + // bare username (browser /agent-api gets this from /api/auth/me via nginx). remoteUser = profile.startsWith("u-") ? profile.substring(2) : profile; postJson("/api/profile/switch", Map.of("name", profile)); } diff --git a/backend/src/main/java/com/dbaagent/service/AuthSessionService.java b/backend/src/main/java/com/dbaagent/service/AuthSessionService.java index 1e40548..3effbf0 100644 --- a/backend/src/main/java/com/dbaagent/service/AuthSessionService.java +++ b/backend/src/main/java/com/dbaagent/service/AuthSessionService.java @@ -155,6 +155,32 @@ public void writeSessionCookies(HttpServletResponse response, SessionAuthenticat response.addHeader(HttpHeaders.SET_COOKIE, buildRefreshCookie(sessionAuthentication.refreshToken()).toString()); } + /** + * Rewrite the access cookie in place (same session id) so View as can stamp + * or clear {@code impUid} without rotating the refresh token. {@code impersonateUserId} + * null clears the claim. + */ + public void reissueAccessToken( + HttpServletResponse response, + String sessionId, + User sessionOwner, + Long impersonateUserId + ) { + if (response == null || sessionId == null || sessionId.isBlank() || sessionOwner == null) { + return; + } + Role role = sessionOwner.getRoleEnum(); + String accessToken = jwtUtil.generateAccessToken( + sessionOwner.getUsername(), + sessionId, + role, + role.getPermissions(), + Duration.ofMinutes(accessMinutes), + impersonateUserId + ); + response.addHeader(HttpHeaders.SET_COOKIE, buildAccessCookie(accessToken).toString()); + } + public void writeImpersonationCookie(HttpServletResponse response, String cookieName, long targetUserId) { response.addHeader(HttpHeaders.SET_COOKIE, ResponseCookie.from(cookieName, Long.toString(targetUserId)) .httpOnly(true) diff --git a/backend/src/main/java/com/dbaagent/service/ChatRetrievalContextService.java b/backend/src/main/java/com/dbaagent/service/ChatRetrievalContextService.java index 340910e..e9e9cbe 100644 --- a/backend/src/main/java/com/dbaagent/service/ChatRetrievalContextService.java +++ b/backend/src/main/java/com/dbaagent/service/ChatRetrievalContextService.java @@ -4,6 +4,7 @@ import com.dbaagent.model.QualifiedTableName; import com.dbaagent.model.SchemaMetadata; import com.dbaagent.model.TrainingDataEmbedding; +import com.dbaagent.service.security.AccessControlService; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -28,6 +29,8 @@ public class ChatRetrievalContextService { private final TrainingService trainingService; private final CompanyKnowledgeService companyKnowledgeService; private final ObjectMapper objectMapper; + private final UserDataAccessPolicyService userDataAccessPolicyService; + private final AccessControlService accessControlService; @Value("${app.chat.rag.general-top-k:20}") private int ragGeneralTopK; @@ -57,6 +60,7 @@ public RetrievedContextResult buildContext( String actualUserQuestion, SchemaMetadata schema ) { + SchemaMetadata scopedSchema = scopeSchema(connectionId, schema); RetrievalIntent retrievalIntent = detectRetrievalIntent(actualUserQuestion); if (isSimpleSchemaQuestion(actualUserQuestion)) { log.debug("Skipping RAG for simple schema question"); @@ -65,7 +69,7 @@ public RetrievedContextResult buildContext( // is also a no-op without focus tables, so the section ends up empty (which // is the right behavior for "list all tables"-style queries). CompanyKnowledgeService.RelevantKnowledgeContext knowledgeContext = - companyKnowledgeService.selectFromRagHits(connectionId, List.of(), Set.of(), schema); + companyKnowledgeService.selectFromRagHits(connectionId, List.of(), Set.of(), scopedSchema); return new RetrievedContextResult( "", knowledgeContext.hintContext(), @@ -85,19 +89,24 @@ public RetrievedContextResult buildContext( long ragStart = System.currentTimeMillis(); int retrievalTopK = resolveRetrievalTopK(retrievalIntent); - List ragResults = trainingService.cachedRetrieveRelevant( - connectionId, actualUserQuestion, retrievalTopK); - ragResults = prioritizeTrainingDataByIntent(ragResults, retrievalIntent); + List ragResults = scopeEmbeddings(connectionId, prioritizeTrainingDataByIntent( + trainingService.cachedRetrieveRelevant(connectionId, actualUserQuestion, retrievalTopK), + retrievalIntent + )); - String dbType = schema != null ? schema.getDbType() : null; - Set resolvedTables = schema == null + String dbType = scopedSchema != null ? scopedSchema.getDbType() : null; + Set resolvedTables = scopedSchema == null ? Set.of() - : trainingService.resolveRelevantTables(schema, actualUserQuestion, ragResults, dbType); + : trainingService.resolveRelevantTables(scopedSchema, actualUserQuestion, ragResults, dbType); - List stage2Results = trainingService.retrieveTargetedByTables( - connectionId, resolvedTables, 100); - List stage3Results = trainingService.deduplicateAgainstTargeted( - ragResults, stage2Results); + List stage2Results = scopeEmbeddings( + connectionId, + trainingService.retrieveTargetedByTables(connectionId, resolvedTables, 100) + ); + List stage3Results = scopeEmbeddings( + connectionId, + trainingService.deduplicateAgainstTargeted(ragResults, stage2Results) + ); List allResults = new ArrayList<>(stage2Results); allResults.addAll(stage3Results); @@ -110,7 +119,7 @@ public RetrievedContextResult buildContext( connectionId, allResults, extractTableNamesFromRag(allResults), - schema + scopedSchema ); Set ragTableNames = new LinkedHashSet<>(extractTableNamesFromRag(allResults)); @@ -153,20 +162,25 @@ public RetrievedContextResult buildScopedContext( return buildContext(connectionId, actualUserQuestion, schema); } + SchemaMetadata scopedSchema = scopeSchema(connectionId, schema); RetrievalIntent retrievalIntent = detectRetrievalIntent(actualUserQuestion); int retrievalTopK = resolveRetrievalTopK(retrievalIntent); - String dbType = schema.getDbType(); - Set resolvedTables = resolveScopedTables(schema, requestedTables, dbType); + String dbType = scopedSchema.getDbType(); + Set resolvedTables = resolveScopedTables(scopedSchema, requestedTables, dbType); String augmentedQuestion = actualUserQuestion + " tables: " + String.join(", ", requestedTables); - List ragResults = trainingService.cachedRetrieveRelevant( - connectionId, augmentedQuestion, retrievalTopK); - ragResults = prioritizeTrainingDataByIntent(ragResults, retrievalIntent); + List ragResults = scopeEmbeddings(connectionId, prioritizeTrainingDataByIntent( + trainingService.cachedRetrieveRelevant(connectionId, augmentedQuestion, retrievalTopK), + retrievalIntent + )); List stage2Results = resolvedTables.isEmpty() ? List.of() - : trainingService.retrieveTargetedByTables(connectionId, resolvedTables, 100); - List stage3Results = trainingService.deduplicateAgainstTargeted(ragResults, stage2Results); + : scopeEmbeddings(connectionId, trainingService.retrieveTargetedByTables(connectionId, resolvedTables, 100)); + List stage3Results = scopeEmbeddings( + connectionId, + trainingService.deduplicateAgainstTargeted(ragResults, stage2Results) + ); List allResults = new ArrayList<>(stage2Results); allResults.addAll(stage3Results); @@ -179,7 +193,7 @@ public RetrievedContextResult buildScopedContext( connectionId, allResults, selectorFocus, - schema + scopedSchema ); Set ragTableNames = new LinkedHashSet<>(requestedTables); @@ -427,4 +441,22 @@ private Set resolveScopedTables( } return resolved; } + + private SchemaMetadata scopeSchema(String connectionId, SchemaMetadata schema) { + return userDataAccessPolicyService.filterSchemaMetadata( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + schema + ); + } + + private List scopeEmbeddings(String connectionId, List embeddings) { + return userDataAccessPolicyService.filterRagEmbeddings( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + embeddings + ); + } } diff --git a/backend/src/main/java/com/dbaagent/service/ChatService.java b/backend/src/main/java/com/dbaagent/service/ChatService.java index 8ed7228..440d735 100644 --- a/backend/src/main/java/com/dbaagent/service/ChatService.java +++ b/backend/src/main/java/com/dbaagent/service/ChatService.java @@ -2649,7 +2649,12 @@ private ChatResponse retiredProcessMessage(String connectionId, String message, } } - SchemaMetadata schema = schemaScannerService.scanSchema(connectionId); + SchemaMetadata schema = scopeSchemaForActor( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + schemaScannerService.scanSchema(connectionId) + ); AgentDecision schemaAwareDecision = earlyAgenticDecision.useAgenticFlow() ? earlyAgenticDecision : agentOrchestrator.previewDecision(agenticEnabled, actualUserQuestion, questionRoute); @@ -2687,7 +2692,12 @@ private ChatResponse retiredProcessMessage(String connectionId, String message, } // 1. Get Schema Context (needed for metadata fast path and schema-aware flows) - SchemaMetadata schema = schemaScannerService.scanSchema(connectionId); + SchemaMetadata schema = scopeSchemaForActor( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + schemaScannerService.scanSchema(connectionId) + ); log.debug("Schema scan completed in {}ms", System.currentTimeMillis() - startTime); // 2. FAST PATH: Check if we can answer directly from cached metadata (no LLM needed) @@ -3620,7 +3630,12 @@ private ChatResponse executePreparedAgenticChatTurn( prepared.effectiveQuestion() ); - SchemaMetadata schema = preloadedMetadataSchema != null ? preloadedMetadataSchema : schemaScannerService.scanSchema(connectionId); + SchemaMetadata schema = scopeSchemaForActor( + connectionId, + actorUsername, + actorIsAdmin, + preloadedMetadataSchema != null ? preloadedMetadataSchema : schemaScannerService.scanSchema(connectionId) + ); ChatResponse vaultFirstMetadataResponse = tryBuildVaultFirstMetadataResponse( connectionId, prepared.chatId(), @@ -3761,7 +3776,12 @@ private ChatResponse tryBuildVaultFirstMetadataResponse( @Nullable private SchemaMetadata tryLoadSchemaForMetadataResponse(String connectionId) { try { - return schemaScannerService.scanSchema(connectionId); + return scopeSchemaForActor( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + schemaScannerService.scanSchema(connectionId) + ); } catch (SQLException e) { log.warn("Failed to load schema for metadata fast path on {}: {}", connectionId, e.getMessage()); return null; @@ -4076,6 +4096,24 @@ private String resolveExecutionActorUsername(@Nullable String explicitUserId) { return accessControlService.getCurrentUsername(); } + private SchemaMetadata scopeSchemaForActor( + String connectionId, + String actorUsername, + boolean actorIsAdmin, + SchemaMetadata schema + ) { + if (schema == null) { + return null; + } + SchemaMetadata scoped = userDataAccessPolicyService.filterSchemaMetadata( + connectionId, + actorUsername, + actorIsAdmin, + schema + ); + return scoped != null ? scoped : schema; + } + private String currentChatOwnerUsername() { String actorUsername = QueryActorContextHolder.currentUsername(); if (actorUsername != null && !actorUsername.isBlank()) { @@ -4938,7 +4976,12 @@ private List dedupeGuardrails( public RetrievedContextResult debugRagContext(String connectionId, String question) { String actualQuestion = extractActualUserQuestion(question); try { - SchemaMetadata schema = schemaScannerService.scanSchema(connectionId); + SchemaMetadata schema = scopeSchemaForActor( + connectionId, + accessControlService.getCurrentUsername(), + accessControlService.isCurrentUserAdmin(), + schemaScannerService.scanSchema(connectionId) + ); return chatRetrievalContextService.buildContext(connectionId, actualQuestion, schema); } catch (Exception e) { log.warn("Failed to load schema for debug RAG context, using empty schema: {}", e.getMessage()); diff --git a/backend/src/main/java/com/dbaagent/service/ConnectionChatAccessPolicyService.java b/backend/src/main/java/com/dbaagent/service/ConnectionChatAccessPolicyService.java index d3175e4..ed92011 100644 --- a/backend/src/main/java/com/dbaagent/service/ConnectionChatAccessPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/ConnectionChatAccessPolicyService.java @@ -76,6 +76,12 @@ public Optional getPolicyResponse(String con .map(this::toResponse); } + @Transactional(readOnly = true) + public boolean hasActivePolicy(String connectionId) { + return connectionId != null && !connectionId.isBlank() + && policyRepository.existsByConnectionIdAndActiveTrue(connectionId); + } + @Transactional(readOnly = true) public EffectivePolicy resolveEffectivePolicy(String connectionId, String username, boolean actorIsAdmin) { if (actorIsAdmin || username == null || username.isBlank()) { @@ -109,6 +115,8 @@ public ConnectionChatAccessPolicyResponse savePolicy( policy.setBlockedSensitivityCategories(parsedPolicy.blockedSensitivityCategories()); policy.setDeniedTables(parsedPolicy.deniedTables()); policy.setDeniedColumns(parsedPolicy.deniedColumns()); + policy.setAllowedSchemas(parsedPolicy.allowedSchemas()); + policy.setAllowAggregates(false); policy.setBlockMode(parsedPolicy.blockMode()); policy.setRedactMode(parsedPolicy.redactMode()); policy.setActive(active == null || active); @@ -140,6 +148,8 @@ public PolicyPreviewResponse previewPolicy(String connectionId, String plainEngl .blockedSensitivityCategories(parsedPolicy.blockedSensitivityCategories()) .deniedTables(parsedPolicy.deniedTables()) .deniedColumns(parsedPolicy.deniedColumns()) + .allowedSchemas(parsedPolicy.allowedSchemas()) + .allowAggregates(false) .impactedTables(parsedPolicy.impactedTables()) .impactedColumns(parsedPolicy.impactedColumns()) .blockMode(parsedPolicy.blockMode()) @@ -157,6 +167,8 @@ private ConnectionChatAccessPolicyResponse toResponse(ConnectionChatAccessPolicy .blockedSensitivityCategories(policy.getBlockedSensitivityCategories()) .deniedTables(policy.getDeniedTables()) .deniedColumns(policy.getDeniedColumns()) + .allowedSchemas(policy.getAllowedSchemas()) + .allowAggregates(policy.isAllowAggregates()) .blockMode(policy.isBlockMode()) .redactMode(policy.isRedactMode()) .active(policy.isActive()) @@ -170,6 +182,10 @@ private ConnectionChatAccessPolicyResponse toResponse(ConnectionChatAccessPolicy private EffectivePolicy toEffectivePolicy(ConnectionChatAccessPolicy policy) { ParsedPolicy parsedPolicy = parsedFromPolicy(policy); + Set storedSchemas = normalizeSet(policy.getAllowedSchemas()); + Set allowedSchemas = storedSchemas.isEmpty() + ? new LinkedHashSet<>(parsedPolicy.allowedSchemas()) + : storedSchemas; return new EffectivePolicy( true, policy.getConnectionId(), @@ -177,9 +193,10 @@ private EffectivePolicy toEffectivePolicy(ConnectionChatAccessPolicy policy) { new LinkedHashSet<>(policy.getBlockedSensitivityCategories() == null ? List.of() : policy.getBlockedSensitivityCategories()), new LinkedHashSet<>(policy.getDeniedTables() == null ? List.of() : policy.getDeniedTables()), new LinkedHashSet<>(policy.getDeniedColumns() == null ? List.of() : policy.getDeniedColumns()), - new LinkedHashSet<>(parsedPolicy.allowedSchemas()), + allowedSchemas, policy.isBlockMode(), policy.isRedactMode(), + policy.isAllowAggregates(), policy.getPlainEnglishPolicy(), parsedPolicy.impactedTables(), parsedPolicy.impactedColumns() @@ -689,12 +706,13 @@ public record EffectivePolicy( Set allowedSchemas, boolean blockMode, boolean redactMode, + boolean allowAggregates, String plainEnglishPolicy, List impactedTables, List impactedColumns ) { public static EffectivePolicy none() { - return new EffectivePolicy(false, null, null, Set.of(), Set.of(), Set.of(), Set.of(), false, false, null, List.of(), List.of()); + return new EffectivePolicy(false, null, null, Set.of(), Set.of(), Set.of(), Set.of(), false, false, false, null, List.of(), List.of()); } public boolean protectsAnything() { diff --git a/backend/src/main/java/com/dbaagent/service/ImpersonationService.java b/backend/src/main/java/com/dbaagent/service/ImpersonationService.java index 2d9750a..dc3e7bd 100644 --- a/backend/src/main/java/com/dbaagent/service/ImpersonationService.java +++ b/backend/src/main/java/com/dbaagent/service/ImpersonationService.java @@ -7,6 +7,7 @@ import com.dbaagent.repository.UserRepository; import com.dbaagent.security.CustomUserDetailsService; import com.dbaagent.security.ImpersonationContext; +import com.dbaagent.security.JwtUtil; import jakarta.servlet.http.Cookie; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; @@ -33,10 +34,18 @@ import static org.springframework.http.HttpStatus.NOT_FOUND; /** - * Admin-only profile switch. The admin session (JWT cookies) is unchanged; - * a separate httpOnly cookie names the user to evaluate as. The JWT filter - * overlays that principal onto the SecurityContext for every request except - * the impersonation control plane, logout, and session refresh. + * Admin-only profile switch so an administrator can verify another user's + * connection ACLs and chat/editor policies. + * + *

The admin session stays the real session (logout/refresh/control-plane). + * Identity for policy is the target user: an httpOnly {@code impersonate_user} + * cookie plus an {@code impUid} claim on the access JWT. The JWT filter overlays + * that principal onto the SecurityContext for every request except the + * impersonation control plane, logout, and session refresh. + * + *

The {@code impUid} claim matters for callers that send the access token as + * a Bearer credential without cookies (Agent MCP fallback). Cookie-only overlay + * left those paths running as the administrator, which skipped policy. */ @Service @RequiredArgsConstructor @@ -49,6 +58,7 @@ public class ImpersonationService { private final CustomUserDetailsService userDetailsService; private final AuthSessionService authSessionService; private final SecurityEventService securityEventService; + private final JwtUtil jwtUtil; @Value("${security.auth.enabled:true}") private boolean authEnabled; @@ -56,6 +66,9 @@ public class ImpersonationService { @Value("${security.cookie.impersonate-name:" + DEFAULT_COOKIE_NAME + "}") private String impersonateCookieName; + @Value("${security.cookie.name:auth_token}") + private String accessCookieName; + public ImpersonationContext.State start( User actor, Long targetUserId, @@ -65,6 +78,7 @@ public ImpersonationContext.State start( requireAdminActor(actor); User target = requireAllowedTarget(actor, targetUserId); authSessionService.writeImpersonationCookie(response, impersonateCookieName, target.getId()); + rewriteAccessToken(request, response, actor, target.getId()); securityEventService.log(SecurityEventService.EventRequest.builder() .eventType(SecurityEventType.IMPERSONATION_STARTED) .outcome(SecurityEventOutcome.SUCCESS) @@ -91,6 +105,7 @@ public User stop( requireAdminActor(actor); Optional target = readTargetUser(request); authSessionService.clearImpersonationCookie(response, impersonateCookieName); + rewriteAccessToken(request, response, actor, null); ImpersonationContext.clear(); target.ifPresent(stopped -> securityEventService.log(SecurityEventService.EventRequest.builder() .eventType(SecurityEventType.IMPERSONATION_STOPPED) @@ -145,6 +160,10 @@ public void decorateAuthPayload(HttpServletRequest request, User sessionUser, Ma * Overlay the target principal when the admin JWT (or the auth-disabled * synthetic admin) is already in the SecurityContext. No-ops on the * impersonation control-plane, logout/refresh, MCP tokens, and invalid cookies. + * + *

The target is taken from the {@code impersonate_user} cookie first, then + * from the access token's {@code impUid} claim so Bearer callers without + * cookies still evaluate policy as the viewed-as user. */ public void applyToRequest(HttpServletRequest request) { if (!shouldApply(request)) { @@ -218,6 +237,14 @@ private Optional readTargetUser(HttpServletRequest request) { } private Long readTargetUserId(HttpServletRequest request) { + Long fromCookie = readTargetUserIdFromCookie(request); + if (fromCookie != null) { + return fromCookie; + } + return readTargetUserIdFromJwt(request); + } + + private Long readTargetUserIdFromCookie(HttpServletRequest request) { Cookie[] cookies = request.getCookies(); if (cookies == null) { return null; @@ -230,6 +257,69 @@ private Long readTargetUserId(HttpServletRequest request) { return null; } + private Long readTargetUserIdFromJwt(HttpServletRequest request) { + String jwt = readAccessJwt(request); + if (jwt == null) { + return null; + } + try { + return jwtUtil.extractImpersonateUserId(jwt); + } catch (Exception e) { + return null; + } + } + + private void rewriteAccessToken( + HttpServletRequest request, + HttpServletResponse response, + User sessionOwner, + Long impersonateUserId + ) { + String sessionId = sessionIdFrom(request); + if (sessionId == null) { + return; + } + authSessionService.reissueAccessToken(response, sessionId, sessionOwner, impersonateUserId); + } + + private String sessionIdFrom(HttpServletRequest request) { + Object attr = request.getAttribute("auth.sessionId"); + if (attr instanceof String sid && !sid.isBlank()) { + return sid; + } + String jwt = readAccessJwt(request); + if (jwt == null) { + return null; + } + try { + String sessionId = jwtUtil.extractSessionId(jwt); + return sessionId == null || sessionId.isBlank() ? null : sessionId; + } catch (Exception e) { + return null; + } + } + + private String readAccessJwt(HttpServletRequest request) { + String authorization = request.getHeader(HttpHeaders.AUTHORIZATION); + if (authorization != null && authorization.startsWith("Bearer ")) { + String token = authorization.substring(7); + if (token.startsWith(McpTokenService.TOKEN_PREFIX)) { + return null; + } + return token.isBlank() ? null : token; + } + Cookie[] cookies = request.getCookies(); + if (cookies == null) { + return null; + } + for (Cookie cookie : cookies) { + if (accessCookieName.equals(cookie.getName())) { + return cookie.getValue(); + } + } + return null; + } + private Long parseUserId(String value) { if (value == null || value.isBlank()) { return null; diff --git a/backend/src/main/java/com/dbaagent/service/QueryExecutionContext.java b/backend/src/main/java/com/dbaagent/service/QueryExecutionContext.java index f28e71b..c8ad685 100644 --- a/backend/src/main/java/com/dbaagent/service/QueryExecutionContext.java +++ b/backend/src/main/java/com/dbaagent/service/QueryExecutionContext.java @@ -1,6 +1,8 @@ package com.dbaagent.service; import com.dbaagent.model.QueryExecutionOrigin; +import org.springframework.security.core.Authentication; +import org.springframework.security.core.context.SecurityContextHolder; public record QueryExecutionContext( QueryExecutionOrigin origin, @@ -19,7 +21,7 @@ public static QueryExecutionContext chat() { return new QueryExecutionContext( QueryExecutionOrigin.CHAT, MutationMode.READ_ONLY_ONLY, - QueryActorContextHolder.currentUsername(), + resolveActorUsername(), false, false ); @@ -46,11 +48,15 @@ public static QueryExecutionContext internal() { } public static QueryExecutionContext mcp(String actorUsername) { + return mcp(actorUsername, false); + } + + public static QueryExecutionContext mcp(String actorUsername, boolean actorIsAdmin) { return new QueryExecutionContext( QueryExecutionOrigin.MCP, MutationMode.READ_ONLY_ONLY, actorUsername, - false, + actorIsAdmin, false ); } @@ -66,12 +72,32 @@ public static QueryExecutionContext scheduled() { } public static QueryExecutionContext api(String actorUsername) { + return api(actorUsername, false); + } + + public static QueryExecutionContext api(String actorUsername, boolean actorIsAdmin) { return new QueryExecutionContext( QueryExecutionOrigin.API, MutationMode.READ_ONLY_ONLY, actorUsername, - false, + actorIsAdmin, false ); } + + private static String resolveActorUsername() { + String fromHolder = QueryActorContextHolder.currentUsername(); + if (fromHolder != null && !fromHolder.isBlank()) { + return fromHolder; + } + Authentication authentication = SecurityContextHolder.getContext().getAuthentication(); + if (authentication == null || !authentication.isAuthenticated()) { + return null; + } + String name = authentication.getName(); + if (name == null || name.isBlank() || "anonymousUser".equals(name)) { + return null; + } + return name; + } } diff --git a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java index a87b214..db5db37 100644 --- a/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java +++ b/backend/src/main/java/com/dbaagent/service/UserDataAccessPolicyService.java @@ -1,22 +1,34 @@ package com.dbaagent.service; +import com.dbaagent.model.DatabaseObject; +import com.dbaagent.model.InferredTableRelationship; +import com.dbaagent.model.QueryExecutionOrigin; import com.dbaagent.model.QueryRequest; import com.dbaagent.model.QueryResult; -import com.dbaagent.model.DatabaseObject; import com.dbaagent.model.SchemaMetadata; -import com.dbaagent.model.TableMetadata; import com.dbaagent.model.SecurityEventOutcome; import com.dbaagent.model.SecurityEventType; +import com.dbaagent.model.TableMetadata; +import com.dbaagent.model.TrainingDataEmbedding; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import net.sf.jsqlparser.expression.AnalyticExpression; import net.sf.jsqlparser.expression.BinaryExpression; +import net.sf.jsqlparser.expression.CaseExpression; +import net.sf.jsqlparser.expression.CastExpression; import net.sf.jsqlparser.expression.Expression; -import net.sf.jsqlparser.expression.Parenthesis; import net.sf.jsqlparser.expression.Function; -import net.sf.jsqlparser.expression.operators.relational.ExpressionList; +import net.sf.jsqlparser.expression.LongValue; +import net.sf.jsqlparser.expression.NotExpression; +import net.sf.jsqlparser.expression.Parenthesis; +import net.sf.jsqlparser.expression.SignedExpression; +import net.sf.jsqlparser.expression.WhenClause; import net.sf.jsqlparser.parser.CCJSqlParserUtil; import net.sf.jsqlparser.schema.Column; import net.sf.jsqlparser.schema.Table; +import net.sf.jsqlparser.statement.ExplainStatement; import net.sf.jsqlparser.statement.Statement; import net.sf.jsqlparser.statement.select.AllColumns; import net.sf.jsqlparser.statement.select.AllTableColumns; @@ -48,6 +60,8 @@ public class UserDataAccessPolicyService { ".*\\b(show|list|export|get|give|find|fetch|retrieve|display|tell me|which)\\b.*\\b(email|emails|phone|phones|mobile|bank|credit card|card number|salary|income|ssn|passport|password|token|address)\\b.*", Pattern.CASE_INSENSITIVE ); + private static final Set SUMMARY_AGGREGATES = Set.of("sum", "avg"); + private static final ObjectMapper RAG_METADATA_MAPPER = new ObjectMapper(); private final ConnectionChatAccessPolicyService policyService; private final SecurityEventService securityEventService; @@ -67,13 +81,7 @@ public PromptDecision evaluatePrompt( || mentionsProtectedColumns(normalized, policy) || mentionsProtectedTables(normalized, policy); - boolean requestsAggregateOnly = normalized.contains("how many") - || normalized.contains("count") - || normalized.contains("trend") - || normalized.contains("summary") - || normalized.contains("breakdown"); - - if (policy.blockMode() && requestsSensitiveDetails && !requestsAggregateOnly) { + if (policy.blockMode() && requestsSensitiveDetails) { logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, username, connectionId, "prompt_blocked", Map.of( "message", truncate(message), "blockedSensitivityCategories", policy.blockedSensitivityCategories(), @@ -102,9 +110,21 @@ public QueryGuardDecision enforcePreExecution( QueryRequest queryRequest, QueryExecutionContext executionContext ) { - if (executionContext == null || executionContext.actorUsername() == null || executionContext.actorUsername().isBlank()) { + if (executionContext == null) { + throw new UserDataAccessPolicyException( + "DeepSQL could not determine who is running this query, so it was blocked.", + "POLICY_ACTOR_REQUIRED" + ); + } + if (isActorExempt(executionContext.origin())) { return QueryGuardDecision.allow(ConnectionChatAccessPolicyService.EffectivePolicy.none()); } + if (executionContext.actorUsername() == null || executionContext.actorUsername().isBlank()) { + throw new UserDataAccessPolicyException( + "DeepSQL could not determine who is running this query, so it was blocked.", + "POLICY_ACTOR_REQUIRED" + ); + } ConnectionChatAccessPolicyService.EffectivePolicy policy = policyService.resolveEffectivePolicy( connectionId, @@ -118,48 +138,46 @@ public QueryGuardDecision enforcePreExecution( Map protectedObjects = policyService.buildProtectionDescriptors(policy); + Statement parsed; try { - Statement parsed = CCJSqlParserUtil.parse(queryRequest.getQuery()); - if (parsed instanceof Select select) { - // Enumerate over the WHOLE statement, not just FROM/JOIN of the - // outermost PlainSelect. An allowlist is only sound if the walk is - // total: any node left unvisited is implicitly permitted, which is - // how a subquery, UNION branch, or CTE body reached a schema - // outside the caller's scope. - enforceAllowedSchemas(parsed, policy.allowedSchemas()); - assertProtectedTablesAreInspectable(parsed, collectPlainSelects(select), protectedObjects); - // Likewise inspect every branch. Gating this on getPlainSelect() - // != null skipped protection entirely for a SetOperationList, - // because a UNION's body is not a PlainSelect. - for (PlainSelect branch : collectPlainSelects(select)) { - QueryInspection inspection = inspectPlainSelect(branch, protectedObjects); - if (inspection.selectsWildcardFromProtectedTable || inspection.rawProtectedColumnsSelected) { - logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, executionContext.actorUsername(), connectionId, "sql_blocked", Map.of( - "query", truncate(queryRequest.getQuery()), - "reason", inspection.reason, - "protectedTables", inspection.protectedTables, - "protectedColumns", inspection.protectedColumns - )); - throw new UserDataAccessPolicyException( - "This query would return restricted data for your account, so DeepSQL blocked it before execution.", - "POLICY_SQL_BLOCKED" - ); - } - } - } - } catch (UserDataAccessPolicyException e) { - throw e; + parsed = CCJSqlParserUtil.parse(queryRequest.getQuery()); } catch (Exception e) { - String normalized = queryRequest.getQuery() == null ? "" : queryRequest.getQuery().toLowerCase(Locale.ROOT); - if (containsDangerousProtectedReference(normalized, protectedObjects)) { - logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, executionContext.actorUsername(), connectionId, "sql_blocked_fallback", Map.of( - "query", truncate(queryRequest.getQuery()) + logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, executionContext.actorUsername(), connectionId, "sql_unparseable", Map.of( + "query", truncate(queryRequest.getQuery()) + )); + throw new UserDataAccessPolicyException( + "This query could not be verified against your access policy, so DeepSQL blocked it before execution.", + "POLICY_SQL_UNPARSEABLE" + ); + } + + try { + Statement inspectable = unwrapExplain(parsed); + if (!(inspectable instanceof Select select)) { + throw new UserDataAccessPolicyException( + "This statement shape cannot be verified against your access policy, so DeepSQL blocked it.", + "POLICY_SQL_UNHANDLED" + ); + } + enforceAllowedSchemas(select, policy.allowedSchemas()); + assertProtectedTablesAreInspectable(inspectable, collectPlainSelects(select), protectedObjects); + QueryInspection inspection = inspectSelect(select, protectedObjects, policy.allowAggregates()); + if (inspection.selectsWildcardFromProtectedTable + || inspection.rawProtectedColumnsSelected + || inspection.unresolvedProtectedReference) { + logPolicyEvent(SecurityEventType.CHAT_ACCESS_POLICY_BLOCKED, executionContext.actorUsername(), connectionId, "sql_blocked", Map.of( + "query", truncate(queryRequest.getQuery()), + "reason", inspection.reason, + "protectedTables", inspection.protectedTables, + "protectedColumns", inspection.protectedColumns )); throw new UserDataAccessPolicyException( - "This query appears to target restricted data for your account, so DeepSQL blocked it before execution.", + "This query would return restricted data for your account, so DeepSQL blocked it before execution.", "POLICY_SQL_BLOCKED" ); } + } catch (UserDataAccessPolicyException e) { + throw e; } return QueryGuardDecision.allow(policy); @@ -217,6 +235,57 @@ public SchemaMetadata filterSchemaMetadata( return filtered; } + public List filterRagEmbeddings( + String connectionId, + String username, + boolean actorIsAdmin, + List embeddings + ) { + if (embeddings == null || embeddings.isEmpty()) { + return embeddings == null ? List.of() : embeddings; + } + Set allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin); + if (allowedSchemas.isEmpty()) { + return embeddings; + } + return embeddings.stream() + .filter(embedding -> isRagMetadataInScope(allowedSchemas, embedding == null ? null : embedding.getMetadata())) + .toList(); + } + + public boolean isRagMetadataInScope( + String connectionId, + String username, + boolean actorIsAdmin, + String metadataJson + ) { + Set allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin); + if (allowedSchemas.isEmpty()) { + return true; + } + return isRagMetadataInScope(allowedSchemas, metadataJson); + } + + public List filterInferredRelationships( + String connectionId, + String username, + boolean actorIsAdmin, + List relationships + ) { + if (relationships == null || relationships.isEmpty()) { + return relationships; + } + Set allowedSchemas = allowedSchemasForActor(connectionId, username, actorIsAdmin); + if (allowedSchemas.isEmpty()) { + return relationships; + } + return relationships.stream() + .filter(relationship -> relationship != null + && ConnectionChatAccessPolicyService.isSchemaInScope(schemaFromTableRef(relationship.getSourceTable()), allowedSchemas) + && ConnectionChatAccessPolicyService.isSchemaInScope(schemaFromTableRef(relationship.getTargetTable()), allowedSchemas)) + .toList(); + } + public void assertTableSchemaAllowed( String connectionId, String username, @@ -251,8 +320,11 @@ private String schemaFromTableRef(String tableRef) { return ""; } String normalized = normalizeName(tableRef); - int separator = normalized.lastIndexOf('.'); - return separator > 0 ? normalized.substring(0, separator) : ""; + String[] parts = normalized.split("\\."); + if (parts.length >= 2) { + return parts[parts.length - 2]; + } + return ""; } public QueryResult redactResult( @@ -280,27 +352,18 @@ public QueryResult redactResult( return result; } - Set protectedColumnNames = new LinkedHashSet<>(); - descriptors.values().forEach(descriptor -> { - String qualifiedTable = descriptor.qualifiedTableName(); - descriptor.restrictedColumns().forEach(column -> - protectedColumnNames.add(normalizeName(qualifiedTable + "." + column)) - ); - if (descriptor.protectWholeTable()) { - protectedColumnNames.add(normalizeName(qualifiedTable)); - } - }); + Set redactIndexes = resolveRedactIndexes(result, descriptors); + if (redactIndexes.isEmpty()) { + return result; + } List> redactedRows = new ArrayList<>(); - boolean redacted = false; for (List row : result.getRows()) { List redactedRow = new ArrayList<>(); for (int i = 0; i < row.size(); i++) { - String column = i < result.getColumns().size() ? result.getColumns().get(i) : ""; Object value = row.get(i); - if (shouldRedactColumn(column, protectedColumnNames)) { + if (redactIndexes.contains(i)) { redactedRow.add(redactValue(value)); - redacted = true; } else { redactedRow.add(value); } @@ -308,10 +371,6 @@ public QueryResult redactResult( redactedRows.add(redactedRow); } - if (!redacted) { - return result; - } - QueryResult redactedResult = new QueryResult( result.getColumns(), redactedRows, @@ -329,6 +388,104 @@ public QueryResult redactResult( return redactedResult; } + private Set resolveRedactIndexes( + QueryResult result, + Map descriptors + ) { + Set indexes = new LinkedHashSet<>(); + List> provenance = resolveOutputProvenance(result.getQuery()); + if (provenance != null && provenance.size() == result.getColumns().size()) { + for (int i = 0; i < provenance.size(); i++) { + Set sources = provenance.get(i); + if (sources.isEmpty()) { + continue; + } + for (ColumnReference reference : sources) { + if (isProtectedReference(descriptors, reference)) { + indexes.add(i); + break; + } + } + } + return indexes; + } + + Set protectedColumnNames = new LinkedHashSet<>(); + descriptors.values().forEach(descriptor -> { + String qualifiedTable = descriptor.qualifiedTableName(); + descriptor.restrictedColumns().forEach(column -> + protectedColumnNames.add(normalizeName(qualifiedTable + "." + column)) + ); + if (descriptor.protectWholeTable()) { + protectedColumnNames.add(normalizeName(qualifiedTable)); + } + }); + for (int i = 0; i < result.getColumns().size(); i++) { + if (shouldRedactColumn(result.getColumns().get(i), protectedColumnNames)) { + indexes.add(i); + } + } + return indexes; + } + + private List> resolveOutputProvenance(String sql) { + if (sql == null || sql.isBlank()) { + return null; + } + try { + Statement parsed = unwrapExplain(CCJSqlParserUtil.parse(sql)); + if (!(parsed instanceof Select select)) { + return null; + } + PlainSelect outermost = outermostPlainSelect(select); + if (outermost == null || outermost.getSelectItems() == null) { + return null; + } + Map aliasToTable = buildAliasMap(outermost); + String defaultTableName = resolveDefaultTableName(outermost, aliasToTable); + List> provenance = new ArrayList<>(); + outermost.getSelectItems().forEach(item -> { + Set referenced = new LinkedHashSet<>(); + Expression expression = item.getExpression(); + if (expression instanceof AllColumns || expression instanceof AllTableColumns) { + provenance.clear(); + return; + } + collectColumns(expression, referenced, aliasToTable, defaultTableName); + provenance.add(referenced); + }); + return provenance.size() == outermost.getSelectItems().size() ? provenance : null; + } catch (Exception e) { + return null; + } + } + + private PlainSelect outermostPlainSelect(Select select) { + if (select == null) { + return null; + } + PlainSelect plainSelect = asPlainSelect(select); + if (plainSelect != null) { + return plainSelect; + } + SetOperationList setOperationList = asSetOperationList(select); + if (setOperationList != null && setOperationList.getSelects() != null && !setOperationList.getSelects().isEmpty()) { + return outermostPlainSelect(setOperationList.getSelects().getFirst()); + } + if (select instanceof ParenthesedSelect parenthesedSelect) { + return outermostPlainSelect(parenthesedSelect.getSelect()); + } + return null; + } + + private PlainSelect asPlainSelect(Select select) { + return select instanceof PlainSelect plainSelect ? plainSelect : null; + } + + private SetOperationList asSetOperationList(Select select) { + return select instanceof SetOperationList setOperationList ? setOperationList : null; + } + private boolean mentionsProtectedColumns(String normalized, ConnectionChatAccessPolicyService.EffectivePolicy policy) { return policy.impactedColumns().stream().anyMatch(column -> normalized.contains(column.toLowerCase(Locale.ROOT)) || normalized.contains(column.substring(column.indexOf('.') + 1).toLowerCase(Locale.ROOT))); @@ -338,20 +495,65 @@ private boolean mentionsProtectedTables(String normalized, ConnectionChatAccessP return policy.impactedTables().stream().anyMatch(table -> normalized.contains(table.toLowerCase(Locale.ROOT))); } + private Statement unwrapExplain(Statement parsed) { + if (parsed instanceof ExplainStatement explain && explain.getStatement() != null) { + return unwrapExplain(explain.getStatement()); + } + return parsed; + } + /** - * Fails closed on statement shapes inspection cannot reach. + * Enforces the schema allowlist over every table reference in the statement. * - * TablesNamesFinder sees every table in the statement; collectPlainSelects - * deliberately does not descend into a select nested inside FROM/JOIN/WHERE/ - * HAVING, because enumerating arbitrary expression trees correctly is the very - * thing that went wrong here the first time. So instead of trying harder to - * walk, compare the two: when a protected table is referenced somewhere the - * column inspection could not examine, refuse the query. + * Uses JSqlParser's TablesNamesFinder rather than a hand-rolled walk of + * FROM and JOIN. Enumeration has to be exhaustive by construction: an + * allowlist implemented as a partial walk implicitly permits every syntax + * position the walker forgot (subquery, UNION branch, CTE body). * - * A syntax form we failed to enumerate must never become an implicit permit -- - * that is exactly how a subquery, a UNION branch and a CTE body each evaded - * the schema allowlist. Refusing costs a conservative block on some safe - * nested aggregates; allowing costs the data. + * Unqualified names fail closed when a schema allowlist is present: they + * resolve through search_path and could be any schema. CTE aliases are + * skipped because they are not tables. + */ + private void enforceAllowedSchemas(Select select, Set allowedSchemas) { + if (allowedSchemas == null || allowedSchemas.isEmpty()) { + return; + } + Set cteNames = new LinkedHashSet<>(); + collectCteNames(select, cteNames); + Set tables = findReferencedTables(select); + Set referencedSchemas = new LinkedHashSet<>(); + for (String tableName : tables) { + String normalized = normalizeName(tableName); + String bare = bareTableName(normalized); + if (cteNames.contains(bare) && !normalized.contains(".")) { + continue; + } + String schema = schemaFromTableRef(normalized); + if (schema.isBlank()) { + throw new UserDataAccessPolicyException( + "This query references an unqualified table which is outside your allowed schema scope.", + "POLICY_SCHEMA_BLOCKED" + ); + } + referencedSchemas.add(schema); + } + for (String schema : referencedSchemas) { + if (!allowedSchemas.contains(normalizeName(schema))) { + throw new UserDataAccessPolicyException( + "This query references schema '" + schema + "' which is outside your allowed schema scope.", + "POLICY_SCHEMA_BLOCKED" + ); + } + } + } + + /** + * Fails closed on statement shapes column inspection cannot reach. + * + * TablesNamesFinder sees every table in the statement; collectPlainSelects + * does not descend into a select nested inside FROM/JOIN/WHERE/HAVING. + * When a protected table is referenced somewhere the column inspection + * could not examine, refuse the query. */ private void assertProtectedTablesAreInspectable( Statement statement, @@ -387,20 +589,12 @@ private void assertProtectedTablesAreInspectable( /** * Does a query's table reference name the protected table? * - * Asymmetric on purpose, because the two sides carry different information. - * ConnectionChatAccessPolicyService.qualifyTable() drops the schema when it is - * "public", so a bare PROTECTED name means public. -- it is not unknown. - * A bare REFERENCE in a query is genuinely unknown: it resolves through the - * session search_path and could be any schema. - * - * reference unqualified -> match on bare name. Ambiguous, so block; the - * search_path may well point at the protected table. - * protected public -> the qualified reference must actually say public. - * marts.customer_profiles is a different table, and - * treating it as protected refused every other - * schema's copy -- which this product's own - * multi-schema fixtures (crm/sales/finance/hr) hit. - * both qualified -> exact match. + * Asymmetric on purpose. {@code qualifyTable()} drops the schema when it is + * {@code public}, so a bare protected name means {@code public.
}. + * A bare reference in a query resolves through search_path and could be + * any schema, so it matches on bare name (ambiguous → block). + * A qualified protection must not catch the same table name in another + * schema ({@code marts.customer_profiles} vs {@code public.customer_profiles}). */ private boolean namesMatch(String protectedName, String referencedName) { String protectedNorm = normalizeName(protectedName); @@ -424,7 +618,7 @@ private String bareName(String normalizedName) { : normalizedName; } - /** Tables named directly in this branch's FROM/JOIN -- what inspection actually saw. */ + /** Tables named directly in this branch's FROM/JOIN. */ private void collectDirectTables(PlainSelect select, Set out) { if (select.getFromItem() instanceof Table table) { out.add(normalizeName(table.getFullyQualifiedName())); @@ -440,9 +634,7 @@ private void collectDirectTables(PlainSelect select, Set out) { /** * Collects every PlainSelect in a statement: the top level, each branch of a - * set operation (UNION/INTERSECT/EXCEPT), parenthesised selects, and every - * CTE body. Callers that inspect only the outermost select leave the rest - * unprotected. + * set operation, parenthesised selects, and every CTE body. */ private List collectPlainSelects(Select select) { List found = new ArrayList<>(); @@ -474,68 +666,90 @@ private void collectPlainSelects(Select select, List found) { } } - /** - * Enforces the schema allowlist over every table reference in the statement. - * - * Uses JSqlParser's TablesNamesFinder rather than a hand-rolled walk of - * FROM and JOIN. The distinction is the whole fix: enumeration has to be - * exhaustive by construction, because an allowlist implemented as a partial - * walk implicitly permits every syntax position the walker forgot -- here a - * subquery in WHERE/HAVING/SELECT, a UNION branch, and a CTE body. - * - * Bare (unqualified) names stay unchecked, exactly as before: they resolve - * through the session search_path, and CTE names are not tables at all. - */ - private void enforceAllowedSchemas(Statement statement, Set allowedSchemas) { - if (allowedSchemas == null || allowedSchemas.isEmpty()) { + private Set findReferencedTables(Select select) { + TablesNamesFinder finder = new TablesNamesFinder(); + List tableList = finder.getTableList((Statement) select); + return tableList == null ? Set.of() : new LinkedHashSet<>(tableList); + } + + private void collectCteNames(Select select, Set names) { + if (select == null) { return; } - Set referencedSchemas = new LinkedHashSet<>(); - for (String qualifiedName : new TablesNamesFinder<>().getTables(statement)) { - if (qualifiedName == null) { - continue; - } - String[] parts = qualifiedName.split("\\."); - if (parts.length >= 2) { - // db.schema.table and schema.table both put the schema second-to-last. - referencedSchemas.add(parts[parts.length - 2]); + if (select.getWithItemsList() != null) { + for (WithItem withItem : select.getWithItemsList()) { + if (withItem.getAlias() != null && withItem.getAlias().getName() != null) { + names.add(normalizeName(withItem.getAlias().getName())); + } + if (withItem.getSelect() != null) { + collectCteNames(withItem.getSelect(), names); + } } } - for (String schema : referencedSchemas) { - if (!allowedSchemas.contains(normalizeName(schema))) { - throw new UserDataAccessPolicyException( - "This query references schema '" + schema + "' which is outside your allowed schema scope.", - "POLICY_SCHEMA_BLOCKED" - ); + SetOperationList setOperationList = asSetOperationList(select); + if (setOperationList != null && setOperationList.getSelects() != null) { + for (Select part : setOperationList.getSelects()) { + collectCteNames(part, names); } } } - private boolean containsDangerousProtectedReference( - String normalizedSql, - Map protectedObjects + private QueryInspection inspectSelect( + Select select, + Map protectedObjects, + boolean allowAggregates ) { - for (ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor : protectedObjects.values()) { - String qualified = normalizeName(descriptor.qualifiedTableName()); - if (descriptor.protectWholeTable() && normalizedSql.contains(qualified)) { - return true; - } - for (String column : descriptor.restrictedColumns()) { - if (normalizedSql.contains(qualified + "." + normalizeName(column))) { - return true; + QueryInspection inspection = new QueryInspection(); + if (select == null) { + inspection.unresolvedProtectedReference = true; + inspection.reason = "Unhandled SELECT shape"; + return inspection; + } + if (select.getWithItemsList() != null) { + for (WithItem withItem : select.getWithItemsList()) { + if (withItem.getSelect() != null) { + inspection.merge(inspectSelect(withItem.getSelect(), protectedObjects, allowAggregates)); } } } - return false; + PlainSelect plainSelect = asPlainSelect(select); + if (plainSelect != null) { + inspection.merge(inspectPlainSelect(plainSelect, protectedObjects, allowAggregates)); + return inspection; + } + SetOperationList setOperationList = asSetOperationList(select); + if (setOperationList != null && setOperationList.getSelects() != null) { + for (Select part : setOperationList.getSelects()) { + inspection.merge(inspectSelect(part, protectedObjects, allowAggregates)); + } + return inspection; + } + if (select instanceof ParenthesedSelect parenthesedSelect && parenthesedSelect.getSelect() != null) { + inspection.merge(inspectSelect(parenthesedSelect.getSelect(), protectedObjects, allowAggregates)); + return inspection; + } + inspection.unresolvedProtectedReference = true; + inspection.reason = "Unhandled SELECT shape"; + return inspection; } private QueryInspection inspectPlainSelect( PlainSelect select, - Map protectedObjects + Map protectedObjects, + boolean allowAggregates ) { Map aliasToTable = buildAliasMap(select); String defaultTableName = resolveDefaultTableName(select, aliasToTable); QueryInspection inspection = new QueryInspection(); + boolean grouped = select.getGroupBy() != null; + inspectFromItem(select.getFromItem(), protectedObjects, allowAggregates, inspection); + if (select.getJoins() != null) { + for (Join join : select.getJoins()) { + inspectFromItem(join.getRightItem(), protectedObjects, allowAggregates, inspection); + } + } + inspectExpressionTree(select.getWhere(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + inspectExpressionTree(select.getHaving(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); if (select.getSelectItems() == null) { return inspection; } @@ -555,7 +769,7 @@ private QueryInspection inspectPlainSelect( if (expression instanceof AllTableColumns allTableColumns) { String tableName = resolveQualifiedTableName(allTableColumns.getTable(), aliasToTable); ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor = lookupDescriptor(protectedObjects, tableName); - if (descriptor != null) { + if (descriptor != null && (descriptor.protectWholeTable() || !descriptor.restrictedColumns().isEmpty())) { inspection.selectsWildcardFromProtectedTable = true; inspection.reason = "SELECT table.* touches protected table"; inspection.protectedTables.add(descriptor.qualifiedTableName()); @@ -565,18 +779,34 @@ private QueryInspection inspectPlainSelect( Set referencedColumns = new LinkedHashSet<>(); collectColumns(expression, referencedColumns, aliasToTable, defaultTableName); - boolean aggregateExpression = containsAggregate(expression); + if (containsSubselect(expression)) { + inspectExpressionTree(expression, aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + } + boolean aggregateExemption = !grouped && isCountStarOnly(expression); + if (allowAggregates && !grouped && isSummaryAggregate(expression)) { + aggregateExemption = true; + } + if (referencedColumns.isEmpty() && !isLiteralOrCountStar(expression) && mentionsAnyProtectedName(expression, protectedObjects)) { + inspection.unresolvedProtectedReference = true; + inspection.reason = "Protected column provenance could not be established"; + return; + } for (ColumnReference reference : referencedColumns) { - ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor = lookupDescriptor(protectedObjects, reference.tableName()); - if (descriptor == null) { + // Outer SELECT id FROM (subquery) t has no concrete FROM table. Nested + // selects are already walked by inspectFromItem; unqualified nested + // table names are refused by assertProtectedTablesAreInspectable + namesMatch. + // Treating blank provenance as unresolved would also block a same-named + // table in another schema (marts.customer_profiles vs public.customer_profiles). + if ((reference.tableName() == null || reference.tableName().isBlank()) && defaultTableName.isBlank()) { continue; } - boolean protectedColumn = descriptor.protectWholeTable() - || descriptor.restrictedColumns().stream().anyMatch(column -> normalizeName(column).equals(normalizeName(reference.columnName()))); - if (protectedColumn) { - inspection.protectedTables.add(descriptor.qualifiedTableName()); - inspection.protectedColumns.add(descriptor.qualifiedTableName() + "." + reference.columnName()); - if (!aggregateExpression) { + if (isProtectedReference(protectedObjects, reference)) { + ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor = lookupDescriptor(protectedObjects, reference.tableName()); + if (descriptor != null) { + inspection.protectedTables.add(descriptor.qualifiedTableName()); + inspection.protectedColumns.add(descriptor.qualifiedTableName() + "." + reference.columnName()); + } + if (!aggregateExemption) { inspection.rawProtectedColumnsSelected = true; inspection.reason = "Raw protected column selected"; } @@ -586,6 +816,66 @@ private QueryInspection inspectPlainSelect( return inspection; } + private void inspectFromItem( + FromItem fromItem, + Map protectedObjects, + boolean allowAggregates, + QueryInspection inspection + ) { + if (fromItem instanceof ParenthesedSelect parenthesedSelect) { + inspection.merge(inspectSelect(parenthesedSelect, protectedObjects, allowAggregates)); + } else if (fromItem instanceof Select select) { + inspection.merge(inspectSelect(select, protectedObjects, allowAggregates)); + } + } + + private void inspectExpressionTree( + Expression expression, + Map aliasToTable, + String defaultTableName, + Map protectedObjects, + boolean allowAggregates, + QueryInspection inspection + ) { + if (expression instanceof ParenthesedSelect parenthesedSelect) { + inspection.merge(inspectSelect(parenthesedSelect, protectedObjects, allowAggregates)); + return; + } + if (expression instanceof Select select) { + inspection.merge(inspectSelect(select, protectedObjects, allowAggregates)); + return; + } + if (expression instanceof Function function && function.getParameters() != null) { + for (Expression parameter : function.getParameters().getExpressions()) { + inspectExpressionTree(parameter, aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + } + return; + } + if (expression instanceof BinaryExpression binaryExpression) { + inspectExpressionTree(binaryExpression.getLeftExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + inspectExpressionTree(binaryExpression.getRightExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + return; + } + if (expression instanceof Parenthesis parenthesis) { + inspectExpressionTree(parenthesis.getExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + return; + } + if (expression instanceof CastExpression castExpression) { + inspectExpressionTree(castExpression.getLeftExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + return; + } + if (expression instanceof CaseExpression caseExpression) { + inspectExpressionTree(caseExpression.getSwitchExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + if (caseExpression.getWhenClauses() != null) { + for (WhenClause whenClause : caseExpression.getWhenClauses()) { + inspectExpressionTree(whenClause.getWhenExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + inspectExpressionTree(whenClause.getThenExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + } + } + inspectExpressionTree(caseExpression.getElseExpression(), aliasToTable, defaultTableName, protectedObjects, allowAggregates, inspection); + } + } + private String resolveDefaultTableName(PlainSelect select, Map aliasToTable) { if (select == null) { return ""; @@ -659,6 +949,10 @@ private void collectColumns( } return; } + if (expression instanceof AnalyticExpression analyticExpression) { + collectColumns(analyticExpression.getExpression(), references, aliasToTable, defaultTableName); + return; + } if (expression instanceof BinaryExpression binaryExpression) { collectColumns(binaryExpression.getLeftExpression(), references, aliasToTable, defaultTableName); collectColumns(binaryExpression.getRightExpression(), references, aliasToTable, defaultTableName); @@ -666,6 +960,29 @@ private void collectColumns( } if (expression instanceof Parenthesis parenthesis) { collectColumns(parenthesis.getExpression(), references, aliasToTable, defaultTableName); + return; + } + if (expression instanceof CastExpression castExpression) { + collectColumns(castExpression.getLeftExpression(), references, aliasToTable, defaultTableName); + return; + } + if (expression instanceof SignedExpression signedExpression) { + collectColumns(signedExpression.getExpression(), references, aliasToTable, defaultTableName); + return; + } + if (expression instanceof NotExpression notExpression) { + collectColumns(notExpression.getExpression(), references, aliasToTable, defaultTableName); + return; + } + if (expression instanceof CaseExpression caseExpression) { + collectColumns(caseExpression.getSwitchExpression(), references, aliasToTable, defaultTableName); + if (caseExpression.getWhenClauses() != null) { + for (WhenClause whenClause : caseExpression.getWhenClauses()) { + collectColumns(whenClause.getWhenExpression(), references, aliasToTable, defaultTableName); + collectColumns(whenClause.getThenExpression(), references, aliasToTable, defaultTableName); + } + } + collectColumns(caseExpression.getElseExpression(), references, aliasToTable, defaultTableName); } } @@ -701,16 +1018,73 @@ private ConnectionChatAccessPolicyService.ProtectionDescriptor lookupDescriptor( return null; } - private boolean containsAggregate(Expression expression) { - if (expression instanceof Function function) { - String name = function.getName(); - return name != null && List.of("count", "sum", "avg", "min", "max", "array_agg", "string_agg").contains(name.toLowerCase(Locale.ROOT)); + private boolean isProtectedReference( + Map protectedObjects, + ColumnReference reference + ) { + ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor = lookupDescriptor(protectedObjects, reference.tableName()); + if (descriptor == null) { + return false; } - if (expression instanceof BinaryExpression binaryExpression) { - return containsAggregate(binaryExpression.getLeftExpression()) || containsAggregate(binaryExpression.getRightExpression()); + return descriptor.protectWholeTable() + || descriptor.restrictedColumns().stream().anyMatch(column -> normalizeName(column).equals(normalizeName(reference.columnName()))); + } + + private boolean isCountStarOnly(Expression expression) { + if (!(expression instanceof Function function)) { + return false; } - if (expression instanceof Parenthesis parenthesis) { - return containsAggregate(parenthesis.getExpression()); + if (function.getName() == null || !"count".equalsIgnoreCase(function.getName())) { + return false; + } + if (function.isAllColumns()) { + return true; + } + if (function.getParameters() == null || function.getParameters().getExpressions() == null + || function.getParameters().getExpressions().isEmpty()) { + return true; + } + if (function.getParameters().getExpressions().size() != 1) { + return false; + } + Expression parameter = function.getParameters().getExpressions().getFirst(); + if (parameter instanceof AllColumns) { + return true; + } + return parameter instanceof LongValue longValue && longValue.getValue() == 1; + } + + private boolean isSummaryAggregate(Expression expression) { + if (!(expression instanceof Function function) || function.getName() == null) { + return false; + } + return SUMMARY_AGGREGATES.contains(function.getName().toLowerCase(Locale.ROOT)); + } + + private boolean isLiteralOrCountStar(Expression expression) { + return isCountStarOnly(expression) + || expression instanceof LongValue + || (expression != null && expression.getClass().getSimpleName().contains("Value")); + } + + private boolean containsSubselect(Expression expression) { + return expression instanceof Select || expression instanceof ParenthesedSelect; + } + + private boolean mentionsAnyProtectedName( + Expression expression, + Map protectedObjects + ) { + if (expression == null) { + return false; + } + String rendered = normalizeName(expression.toString()); + for (ConnectionChatAccessPolicyService.ProtectionDescriptor descriptor : protectedObjects.values()) { + for (String column : descriptor.restrictedColumns()) { + if (rendered.contains(normalizeName(column))) { + return true; + } + } } return false; } @@ -736,6 +1110,80 @@ private Object redactValue(Object value) { return "[redacted:" + rendered.length() + "]"; } + private boolean isRagMetadataInScope(Set allowedSchemas, String metadataJson) { + Set tableRefs = extractTableRefsFromMetadata(metadataJson); + if (tableRefs.isEmpty()) { + return false; + } + for (String tableRef : tableRefs) { + if (!ConnectionChatAccessPolicyService.isSchemaInScope(schemaFromTableRef(tableRef), allowedSchemas)) { + return false; + } + } + return true; + } + + private Set extractTableRefsFromMetadata(String metadataJson) { + Set tables = new LinkedHashSet<>(); + if (metadataJson == null || metadataJson.isBlank()) { + return tables; + } + try { + JsonNode node = RAG_METADATA_MAPPER.readTree(metadataJson); + addTextualTable(tables, node, "tableName"); + addTextualTable(tables, node, "objectName"); + addTextualTable(tables, node, "schema"); + addTextualTable(tables, node, "schemaName"); + if (node.has("tablesUsed") && node.get("tablesUsed").isTextual()) { + for (String table : node.get("tablesUsed").asText("").split(",")) { + addTableRef(tables, table); + } + } + addArrayTables(tables, node.get("linkedTables")); + if (node.has("linkedColumns") && node.get("linkedColumns").isArray()) { + for (JsonNode linkedColumn : node.get("linkedColumns")) { + String columnReference = linkedColumn.asText("").trim(); + int lastDot = columnReference.lastIndexOf('.'); + if (lastDot > 0) { + addTableRef(tables, columnReference.substring(0, lastDot)); + } + } + } + } catch (Exception ignored) { + return Set.of(); + } + return tables; + } + + private void addTextualTable(Set tables, JsonNode node, String field) { + if (node != null && node.has(field) && node.get(field).isTextual()) { + addTableRef(tables, node.get(field).asText()); + } + } + + private void addArrayTables(Set tables, JsonNode array) { + if (array == null || !array.isArray()) { + return; + } + for (JsonNode item : array) { + addTableRef(tables, item.asText()); + } + } + + private void addTableRef(Set tables, String raw) { + if (raw == null) { + return; + } + String trimmed = raw.trim(); + if (!trimmed.isBlank()) { + tables.add(trimmed); + } + } + + private boolean isActorExempt(QueryExecutionOrigin origin) { + return origin == QueryExecutionOrigin.INTERNAL || origin == QueryExecutionOrigin.SCHEDULED; + } + private void logPolicyEvent(SecurityEventType eventType, String username, String connectionId, String reason, Map metadata) { securityEventService.log(SecurityEventService.EventRequest.builder() .eventType(eventType) @@ -758,6 +1206,11 @@ private String normalizeName(String value) { return value == null ? "" : value.trim().replace("\"", "").replace("`", "").toLowerCase(Locale.ROOT); } + private String bareTableName(String qualified) { + int separator = qualified.lastIndexOf('.'); + return separator >= 0 ? qualified.substring(separator + 1) : qualified; + } + public record PromptDecision( boolean allowed, String responseMessage, @@ -784,9 +1237,21 @@ public static QueryGuardDecision allow(ConnectionChatAccessPolicyService.Effecti private static final class QueryInspection { private boolean selectsWildcardFromProtectedTable; private boolean rawProtectedColumnsSelected; + private boolean unresolvedProtectedReference; private String reason; private final List protectedTables = new ArrayList<>(); private final List protectedColumns = new ArrayList<>(); + + private void merge(QueryInspection other) { + this.selectsWildcardFromProtectedTable |= other.selectsWildcardFromProtectedTable; + this.rawProtectedColumnsSelected |= other.rawProtectedColumnsSelected; + this.unresolvedProtectedReference |= other.unresolvedProtectedReference; + if (this.reason == null) { + this.reason = other.reason; + } + this.protectedTables.addAll(other.protectedTables); + this.protectedColumns.addAll(other.protectedColumns); + } } private record ColumnReference(String tableName, String columnName) { diff --git a/backend/src/main/resources/db/migration/V115__chat_access_policy_allowed_schemas.sql b/backend/src/main/resources/db/migration/V115__chat_access_policy_allowed_schemas.sql new file mode 100644 index 0000000..0b10f97 --- /dev/null +++ b/backend/src/main/resources/db/migration/V115__chat_access_policy_allowed_schemas.sql @@ -0,0 +1,3 @@ +ALTER TABLE connection_chat_access_policy + ADD COLUMN IF NOT EXISTS allowed_schemas jsonb, + ADD COLUMN IF NOT EXISTS allow_aggregates BOOLEAN NOT NULL DEFAULT FALSE; diff --git a/backend/src/test/java/com/dbaagent/security/JwtUtilImpersonationClaimTest.java b/backend/src/test/java/com/dbaagent/security/JwtUtilImpersonationClaimTest.java new file mode 100644 index 0000000..ffbc81c --- /dev/null +++ b/backend/src/test/java/com/dbaagent/security/JwtUtilImpersonationClaimTest.java @@ -0,0 +1,58 @@ +package com.dbaagent.security; + +import com.dbaagent.model.Permission; +import com.dbaagent.model.Role; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.test.util.ReflectionTestUtils; + +import java.time.Duration; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +class JwtUtilImpersonationClaimTest { + + private JwtUtil jwtUtil; + + @BeforeEach + void setUp() { + jwtUtil = new JwtUtil(); + ReflectionTestUtils.setField(jwtUtil, "jwtSecret", "x".repeat(32)); + ReflectionTestUtils.setField(jwtUtil, "authEnabled", true); + ReflectionTestUtils.setField(jwtUtil, "activeProfiles", "dev"); + ReflectionTestUtils.setField(jwtUtil, "accessTokenMinutes", 15L); + jwtUtil.initialize(); + } + + @Test + void accessTokenRoundTripsImpersonateUserId() { + String token = jwtUtil.generateAccessToken( + "admin", + "sess-1", + Role.ADMIN, + Set.of(Permission.MANAGE_USERS), + Duration.ofMinutes(15), + 2L + ); + + assertEquals("admin", jwtUtil.extractUsername(token)); + assertEquals("sess-1", jwtUtil.extractSessionId(token)); + assertEquals(2L, jwtUtil.extractImpersonateUserId(token)); + assertEquals("ADMIN", jwtUtil.extractRole(token)); + } + + @Test + void accessTokenWithoutClaimHasNoImpersonateUserId() { + String token = jwtUtil.generateAccessToken( + "admin", + "sess-1", + Role.ADMIN, + Set.of(Permission.MANAGE_USERS), + Duration.ofMinutes(15) + ); + + assertNull(jwtUtil.extractImpersonateUserId(token)); + } +} diff --git a/backend/src/test/java/com/dbaagent/service/AgentBridgeServiceImpersonationTest.java b/backend/src/test/java/com/dbaagent/service/AgentBridgeServiceImpersonationTest.java new file mode 100644 index 0000000..ed90240 --- /dev/null +++ b/backend/src/test/java/com/dbaagent/service/AgentBridgeServiceImpersonationTest.java @@ -0,0 +1,88 @@ +package com.dbaagent.service; + +import com.dbaagent.model.User; +import com.dbaagent.security.ImpersonationContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.test.util.ReflectionTestUtils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class AgentBridgeServiceImpersonationTest { + + @Mock + private McpTokenService mcpTokenService; + + private AgentBridgeService agentBridgeService; + + @BeforeEach + void setUp() { + agentBridgeService = new AgentBridgeService(mcpTokenService); + ReflectionTestUtils.setField(agentBridgeService, "provisionEnabled", true); + ReflectionTestUtils.setField(agentBridgeService, "provisionSecret", "secret"); + ReflectionTestUtils.setField(agentBridgeService, "provisionerUrl", "http://127.0.0.1:9/provision"); + ReflectionTestUtils.setField(agentBridgeService, "sessionWindowDays", 7L); + } + + @AfterEach + void tearDown() { + ImpersonationContext.clear(); + } + + @Test + void ensureProfileDoesNotFallBackToAdminSessionTokenWhileImpersonating() { + User admin = new User(); + admin.setId(1L); + admin.setUsername("admin"); + admin.setRole("ADMIN"); + User editor = new User(); + editor.setId(2L); + editor.setUsername("marts-editor"); + editor.setRole("DEVELOPER"); + ImpersonationContext.enter(new ImpersonationContext.State(admin, editor)); + + when(mcpTokenService.createTokenForUser(eq("marts-editor"), any(), any())) + .thenThrow(new IllegalStateException("mint failed")); + + AgentBridgeService.ProvisioningException ex = assertThrows( + AgentBridgeService.ProvisioningException.class, + () -> agentBridgeService.ensureProfile("marts-editor", "admin-session-jwt", "conn-1") + ); + assertEquals( + "Could not mint an MCP token for marts-editor while viewing as that user", + ex.getMessage() + ); + } + + @Test + void ensureProfileRefusesDisabledProvisioningWhileImpersonating() { + ReflectionTestUtils.setField(agentBridgeService, "provisionEnabled", false); + User admin = new User(); + admin.setId(1L); + admin.setUsername("admin"); + admin.setRole("ADMIN"); + User editor = new User(); + editor.setId(2L); + editor.setUsername("marts-editor"); + editor.setRole("DEVELOPER"); + ImpersonationContext.enter(new ImpersonationContext.State(admin, editor)); + + AgentBridgeService.ProvisioningException ex = assertThrows( + AgentBridgeService.ProvisioningException.class, + () -> agentBridgeService.ensureProfile("marts-editor", "admin-session-jwt", "conn-1") + ); + assertEquals( + "Agent provisioning is disabled; View as cannot bind a user-scoped Agent token", + ex.getMessage() + ); + } +} diff --git a/backend/src/test/java/com/dbaagent/service/ChatServiceRoutingTest.java b/backend/src/test/java/com/dbaagent/service/ChatServiceRoutingTest.java index 8e9a099..58b3d16 100644 --- a/backend/src/test/java/com/dbaagent/service/ChatServiceRoutingTest.java +++ b/backend/src/test/java/com/dbaagent/service/ChatServiceRoutingTest.java @@ -164,6 +164,8 @@ void setUp() { .thenReturn(UserDataAccessPolicyService.PromptDecision.allow(ConnectionChatAccessPolicyService.EffectivePolicy.none())); lenient().when(userDataAccessPolicyService.decorateQuestionWithPolicy(any(), anyString())) .thenAnswer(invocation -> invocation.getArgument(1, String.class)); + lenient().when(userDataAccessPolicyService.filterSchemaMetadata(any(), any(), anyBoolean(), any())) + .thenAnswer(invocation -> invocation.getArgument(3)); lenient().when(contextAssembler.formatRowCount(anyLong())) .thenAnswer(invocation -> String.valueOf(invocation.getArgument(0, Long.class))); lenient().when(contextAssembler.formatBytes(anyLong())) diff --git a/backend/src/test/java/com/dbaagent/service/ConnectionChatAccessPolicyServiceTest.java b/backend/src/test/java/com/dbaagent/service/ConnectionChatAccessPolicyServiceTest.java index f1ee627..a78aafc 100644 --- a/backend/src/test/java/com/dbaagent/service/ConnectionChatAccessPolicyServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/ConnectionChatAccessPolicyServiceTest.java @@ -115,6 +115,7 @@ void previewPolicy_scopesTypedColumnConstraintsToAllowedSchema() throws SQLExcep "marts.fct_enrollment.currency", "marts.dim_ott_subscription.currency" ); + assertThat(preview.getAllowedSchemas()).containsExactly("marts"); } @Test diff --git a/backend/src/test/java/com/dbaagent/service/ImpersonationServiceTest.java b/backend/src/test/java/com/dbaagent/service/ImpersonationServiceTest.java index 029ea4b..92a94da 100644 --- a/backend/src/test/java/com/dbaagent/service/ImpersonationServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/ImpersonationServiceTest.java @@ -7,6 +7,7 @@ import com.dbaagent.repository.UserRepository; import com.dbaagent.security.CustomUserDetailsService; import com.dbaagent.security.ImpersonationContext; +import com.dbaagent.security.JwtUtil; import jakarta.servlet.http.Cookie; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -54,6 +55,9 @@ class ImpersonationServiceTest { @Mock private SecurityEventService securityEventService; + @Mock + private JwtUtil jwtUtil; + @InjectMocks private ImpersonationService impersonationService; @@ -61,6 +65,7 @@ class ImpersonationServiceTest { void setUp() { ReflectionTestUtils.setField(impersonationService, "authEnabled", true); ReflectionTestUtils.setField(impersonationService, "impersonateCookieName", "impersonate_user"); + ReflectionTestUtils.setField(impersonationService, "accessCookieName", "auth_token"); } @AfterEach @@ -82,6 +87,7 @@ void startWritesCookieAndAudits() { assertEquals("marts-editor", state.targetUsername()); verify(authSessionService).writeImpersonationCookie(response, "impersonate_user", 2L); + verify(authSessionService, never()).reissueAccessToken(any(), any(), any(), any()); ArgumentCaptor captor = ArgumentCaptor.forClass(SecurityEventService.EventRequest.class); verify(securityEventService).log(captor.capture()); @@ -133,6 +139,54 @@ void startRejectsNonAdminActor() { assertEquals(403, ex.getStatusCode().value()); } + @Test + void startRewritesAccessTokenWithImpersonationClaim() { + User admin = user(1L, "admin", "ADMIN"); + User editor = user(2L, "marts-editor", "DEVELOPER"); + when(userRepository.findById(2L)).thenReturn(Optional.of(editor)); + + MockHttpServletRequest request = new MockHttpServletRequest(); + request.setAttribute("auth.sessionId", "sess-1"); + MockHttpServletResponse response = new MockHttpServletResponse(); + + impersonationService.start(admin, 2L, request, response); + + verify(authSessionService).reissueAccessToken(response, "sess-1", admin, 2L); + } + + @Test + void applySwapsPrincipalFromAccessTokenClaimWithoutCookie() { + User admin = user(1L, "admin", "ADMIN"); + User editor = user(2L, "marts-editor", "DEVELOPER"); + when(userRepository.findByUsername("admin")).thenReturn(Optional.of(admin)); + when(userRepository.findById(2L)).thenReturn(Optional.of(editor)); + when(jwtUtil.extractImpersonateUserId("admin.jwt.with.imp")).thenReturn(2L); + UserDetails editorDetails = new org.springframework.security.core.userdetails.User( + "marts-editor", + "x", + List.of(new SimpleGrantedAuthority("ROLE_DEVELOPER"), new SimpleGrantedAuthority("USE_CHAT")) + ); + when(userDetailsService.loadUserByUsername("marts-editor")).thenReturn(editorDetails); + + SecurityContextHolder.getContext().setAuthentication( + new UsernamePasswordAuthenticationToken( + "admin", + null, + List.of(new SimpleGrantedAuthority("ROLE_ADMIN")) + ) + ); + + MockHttpServletRequest request = new MockHttpServletRequest("POST", "/api/connections/abc/query"); + request.setServletPath("/connections/abc/query"); + request.addHeader("Authorization", "Bearer admin.jwt.with.imp"); + + impersonationService.applyToRequest(request); + + assertEquals("marts-editor", SecurityContextHolder.getContext().getAuthentication().getName()); + assertTrue(ImpersonationContext.isActive()); + assertEquals("marts-editor", ImpersonationContext.current().orElseThrow().targetUsername()); + } + @Test void applySwapsPrincipalToTargetUser() { User admin = user(1L, "admin", "ADMIN"); diff --git a/backend/src/test/java/com/dbaagent/service/QueryExecutionContextTest.java b/backend/src/test/java/com/dbaagent/service/QueryExecutionContextTest.java index c7286da..34c23d3 100644 --- a/backend/src/test/java/com/dbaagent/service/QueryExecutionContextTest.java +++ b/backend/src/test/java/com/dbaagent/service/QueryExecutionContextTest.java @@ -17,6 +17,14 @@ void mcpFactoryProducesReadOnlyContextWithMcpOrigin() { assertThat(ctx.mutationConfirmed()).isFalse(); } + @Test + void mcpFactoryHonoursAdminFlagFromSecurityContext() { + QueryExecutionContext ctx = QueryExecutionContext.mcp("admin", true); + assertThat(ctx.origin()).isEqualTo(QueryExecutionOrigin.MCP); + assertThat(ctx.actorUsername()).isEqualTo("admin"); + assertThat(ctx.actorIsAdmin()).isTrue(); + } + @Test void scheduledFactoryProducesMayMutateInternalActor() { QueryExecutionContext ctx = QueryExecutionContext.scheduled(); diff --git a/backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java b/backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java index 9939336..ea5ce8d 100644 --- a/backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/UserDataAccessPolicyServiceTest.java @@ -64,6 +64,21 @@ void evaluatePrompt_blocksDirectSensitiveRequestButAllowsAggregate() { assertThat(allowed.allowed()).isTrue(); } + @Test + void evaluatePrompt_blocksProtectedColumnMentionEvenWhenAskingForACount() { + ConnectionChatAccessPolicyService.EffectivePolicy policy = policy(); + when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy); + + UserDataAccessPolicyService.PromptDecision blocked = service.evaluatePrompt( + "conn-1", + "analyst", + false, + "How many customer emails do we have?" + ); + + assertThat(blocked.allowed()).isFalse(); + } + @Test void enforcePreExecution_blocksRawProtectedColumns() { when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy()); @@ -121,6 +136,7 @@ void enforcePreExecution_blocksQueriesOutsideAllowedSchema() { Set.of("marts"), true, true, + false, "Only schema marts", List.of(), List.of() @@ -144,7 +160,7 @@ private ConnectionChatAccessPolicyService.EffectivePolicy martsOnlyPolicy() { return new ConnectionChatAccessPolicyService.EffectivePolicy( true, "conn-1", "analyst", Set.of(), Set.of(), Set.of(), Set.of("marts"), - true, true, "Only schema marts", List.of(), List.of() + true, true, false, "Only schema marts", List.of(), List.of() ); } @@ -288,6 +304,7 @@ void enforcePreExecution_allowsMartsQueriesWhenOtherSchemasHaveProtectedColumns( Set.of("marts"), true, true, + false, "Only marts; redact amount", List.of(), List.of("marts.fct_enrollment.amount") @@ -317,6 +334,7 @@ void filterDatabaseObjects_keepsOnlyAllowedSchemas() { Set.of("marts"), true, true, + false, "Only schema marts", List.of(), List.of() @@ -344,6 +362,7 @@ void filterSchemaMetadata_dropsOutOfScopeTablesAndRelationships() { Set.of("marts"), true, true, + false, "Only schema marts", List.of(), List.of() @@ -381,6 +400,7 @@ void assertTableSchemaAllowed_blocksOutOfScopeTableMetadata() { Set.of("marts"), true, true, + false, "Only schema marts", List.of(), List.of() @@ -395,6 +415,183 @@ void assertTableSchemaAllowed_blocksOutOfScopeTableMetadata() { service.assertTableSchemaAllowed("conn-1", "analyst", false, "marts.fct_enrollment"); } + @Test + void enforcePreExecution_blocksUnionAndSubqueryOutsideAllowedSchema() { + stubSchemaPolicy(); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("SELECT name FROM marts.fct_enrollment UNION SELECT name FROM sales.customers") + ).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED"); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("WITH x AS (SELECT * FROM crm.customers) SELECT * FROM x") + ).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED"); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("SELECT * FROM (SELECT * FROM sales.customers) t") + ).getErrorCode()).isEqualTo("POLICY_SCHEMA_BLOCKED"); + } + + @Test + void enforcePreExecution_failsClosedOnUnparseableAndUnhandledSql() { + stubSchemaPolicy(); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("SELECT FROM") + ).getErrorCode()).isEqualTo("POLICY_SQL_UNPARSEABLE"); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("ALTER TABLE marts.fct_enrollment RENAME TO fct_enrollment_x") + ).getErrorCode()).isEqualTo("POLICY_SQL_UNHANDLED"); + } + + @Test + void enforcePreExecution_failsClosedWhenActorIsMissingExceptInternalOrigins() { + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> service.enforcePreExecution( + "conn-1", + new QueryRequest("SELECT 1", null, null), + new QueryExecutionContext(QueryExecutionOrigin.MCP, QueryExecutionContext.MutationMode.READ_ONLY_ONLY, null, false, false) + ) + ).getErrorCode()).isEqualTo("POLICY_ACTOR_REQUIRED"); + + assertThat(service.enforcePreExecution( + "conn-1", + new QueryRequest("SELECT 1", null, null), + QueryExecutionContext.internal() + ).policy().present()).isFalse(); + } + + @Test + void enforcePreExecution_countStarIsAllowedButCountOfProtectedColumnIsNot() { + ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = martsAmountPolicy(); + when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy); + when(policyService.buildProtectionDescriptors(schemaPolicy)).thenReturn(Map.of( + "marts.fct_enrollment", + descriptor("marts", "fct_enrollment", false, "amount") + )); + + assertThat(enforce("SELECT COUNT(*) AS n FROM marts.fct_enrollment").policy().allowedSchemas()) + .containsExactly("marts"); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("SELECT COUNT(amount) FROM marts.fct_enrollment") + ).getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED"); + + assertThat(assertThrows( + UserDataAccessPolicyException.class, + () -> enforce("SELECT MIN(amount) FROM marts.fct_enrollment") + ).getErrorCode()).isEqualTo("POLICY_SQL_BLOCKED"); + } + + @Test + void redactResult_usesSourceColumnNotOutputAlias() { + when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(policy()); + + QueryResult aliasHidden = service.redactResult( + "conn-1", + new QueryResult( + List.of("contact"), + List.of(List.of("a@example.com")), + 1, + 1L, + false, + 12L, + "SELECT email AS contact FROM customer_profiles" + ), + chatContext() + ); + assertThat(aliasHidden.getRows()).containsExactly(List.of("[redacted:13]")); + + QueryResult aliasSafe = service.redactResult( + "conn-1", + new QueryResult( + List.of("email"), + List.of(List.of("Hotel One")), + 1, + 1L, + false, + 12L, + "SELECT customer_name AS email FROM customer_profiles" + ), + chatContext() + ); + assertThat(aliasSafe.getRows()).containsExactly(List.of("Hotel One")); + } + + @Test + void filterRagEmbeddings_dropsOutOfScopeDocuments() { + stubSchemaPolicy(); + var inScope = new com.dbaagent.model.TrainingDataEmbedding(); + inScope.setMetadata("{\"tableName\":\"marts.fct_enrollment\"}"); + var outOfScope = new com.dbaagent.model.TrainingDataEmbedding(); + outOfScope.setMetadata("{\"tableName\":\"sales.customers\"}"); + var unknown = new com.dbaagent.model.TrainingDataEmbedding(); + unknown.setMetadata("{}"); + + assertThat(service.filterRagEmbeddings("conn-1", "analyst", false, List.of(inScope, outOfScope, unknown))) + .containsExactly(inScope); + } + + private void stubSchemaPolicy() { + ConnectionChatAccessPolicyService.EffectivePolicy schemaPolicy = new ConnectionChatAccessPolicyService.EffectivePolicy( + true, + "conn-1", + "analyst", + Set.of(), + Set.of(), + Set.of(), + Set.of("marts"), + true, + true, + false, + "Only schema marts", + List.of(), + List.of() + ); + when(policyService.resolveEffectivePolicy("conn-1", "analyst", false)).thenReturn(schemaPolicy); + lenient().when(policyService.buildProtectionDescriptors(schemaPolicy)).thenReturn(Map.of()); + } + + private ConnectionChatAccessPolicyService.EffectivePolicy martsAmountPolicy() { + return new ConnectionChatAccessPolicyService.EffectivePolicy( + true, + "conn-1", + "analyst", + Set.of(), + Set.of(), + Set.of("marts.fct_enrollment.amount"), + Set.of("marts"), + true, + true, + false, + "Only marts; redact amount", + List.of(), + List.of("marts.fct_enrollment.amount") + ); + } + + private UserDataAccessPolicyService.QueryGuardDecision enforce(String sql) { + return service.enforcePreExecution("conn-1", new QueryRequest(sql, null, null), chatContext()); + } + + private QueryExecutionContext chatContext() { + return new QueryExecutionContext( + QueryExecutionOrigin.CHAT, + QueryExecutionContext.MutationMode.READ_ONLY_ONLY, + "analyst", + false, + false + ); + } + private ConnectionChatAccessPolicyService.EffectivePolicy policy() { return new ConnectionChatAccessPolicyService.EffectivePolicy( true, @@ -406,6 +603,7 @@ private ConnectionChatAccessPolicyService.EffectivePolicy policy() { Set.of(), true, true, + false, "No PII", List.of(), List.of("customer_profiles.email", "customer_profiles.phone_number") diff --git a/docker/nginx/default.conf b/docker/nginx/default.conf index b4e2f7d..fe6abd1 100644 --- a/docker/nginx/default.conf +++ b/docker/nginx/default.conf @@ -105,6 +105,7 @@ server { location /agent-api/ { # Require a valid DeepSQL session before reaching the agent. auth_request /__agent_auth; + auth_request_set $deepsql_user $upstream_http_x_remote_user; # Compose service on the internal network. Literal hostname resolves # via Docker DNS. If the agent container is down this route returns @@ -119,9 +120,10 @@ server { proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; proxy_set_header X-Forwarded-Proto $scheme; # Identity for trusted-proxy mode inside the agent container. - # DeepSQL's auth_request already verified the session; the agent - # accepts this header instead of its own login form. - proxy_set_header X-Remote-User admin; + # DeepSQL's auth_request already verified the session; stamp the + # *effective* user from /api/auth/me (View as overlays the target). + # Never hardcode admin — that ran every Agent tab as the admin MCP token. + proxy_set_header X-Remote-User $deepsql_user; proxy_read_timeout 300s; proxy_send_timeout 300s; proxy_connect_timeout 10s; diff --git a/docs/root/CLAUDE.md b/docs/root/CLAUDE.md index fcc4c2f..f46936c 100644 --- a/docs/root/CLAUDE.md +++ b/docs/root/CLAUDE.md @@ -21,9 +21,14 @@ Do NOT ask "should I update CLAUDE.md?" - just update it as part of task complet - Batch multiple related changes into a single commit - Provide a custom commit message if needed -## Recent Changes +- 2026-08-20: Chat access policy enforcement is fail-closed. `UserDataAccessPolicyService` + walks the whole parse tree (`TablesNamesFinder` plus CTEs/UNIONs/subqueries), denies + unparseable or unhandled SQL, requires an actor except `INTERNAL`/`SCHEDULED`, and + exempts only `COUNT(*)` / `COUNT(1)` without `GROUP BY`. MCP and Editor queries take + the actor from `SecurityContext` (MCP tokens included). Allowed schemas are persisted + on the policy row. RAG, brain relationships, and vault-first metadata are schema-scoped. + Public dashboard share is refused on connections with an active policy. -- 2026-08-17: `McpSqlGuardService` (and the matching MCP JS shim) no longer treats `COMMENT` / `CALL` / `REPLACE` as mutating when they appear as table, column, or function names. The guard matches statement verbs: mutating CTEs, `WITH … DELETE`, `FOR UPDATE`, and `EXPLAIN DELETE`. `SELECT * FROM comment` is allowed. Dashboards @@ -822,9 +827,9 @@ though the properties themselves still sit in `application*.properties`. - `DELETE /api/admin/users/{id}` - Delete user (ADMIN only) - `GET /api/admin/roles` - Get all roles with permissions (ADMIN only) - `GET /api/admin/impersonate` - List switchable users and current profile-switch status (ADMIN only) - - `POST /api/admin/impersonate` - `{ userId }` start viewing the product as that user (ADMIN only; cannot target admins or self) + - `POST /api/admin/impersonate` - `{ userId }` start viewing the product as that user (ADMIN only; cannot target admins or self). Sets `impersonate_user` cookie and restamps the access JWT with `impUid` so policy (and Agent Bearer fallback) run as the target while the JWT subject stays the admin. - `DELETE /api/admin/impersonate` - Stop profile switch and restore the admin session - - `GET /api/auth/me` - Get current user's profile including role/permissions; while switching, this is the **target** user plus `impersonating` / `impersonatorUsername` + - `GET /api/auth/me` - Get current user's profile including role/permissions; while switching, this is the **target** user plus `impersonating` / `impersonatorUsername`. Also sets `X-Remote-User` to the effective username for nginx `/agent-api` auth_request. - **Frontend Components**: - `PermissionGuard.jsx` - Wrapper component for permission-based rendering - `UsersTab.jsx` - Admin user management tab in Workspace diff --git a/scripts/local-agent-provisioner.py b/scripts/local-agent-provisioner.py index 6e829fa..3dc59b5 100755 --- a/scripts/local-agent-provisioner.py +++ b/scripts/local-agent-provisioner.py @@ -56,6 +56,7 @@ def write_token_file(home: Path, token: str) -> Path: MCP subprocess (which re-reads this file on every request, see deepsql-phase1-lib.js readTokenFile) never observes a partially-written token mid-rotation. 0600 — same secrecy bar as the profile .env.""" + home.mkdir(parents=True, exist_ok=True) path = token_file_for(home) tmp = path.with_suffix(f".tmp-{os.getpid()}") tmp.write_text((token or "") + "\n") @@ -64,6 +65,40 @@ def write_token_file(home: Path, token: str) -> Path: return path +def iter_profile_homes(hermes_home: Path): + """Yield each ``profiles/`` directory under a Hermes home.""" + profiles = hermes_home / "profiles" + if not profiles.is_dir(): + return + for child in profiles.iterdir(): + if child.is_dir(): + yield child + + +def mirror_token_for_live_mcp(hermes_home: Path, token: str) -> list[Path]: + """Copy the provisioned token onto every file a live MCP process might read. + + Hermes starts one DeepSQL MCP stdio server from whichever profile first + loaded ``mcp_servers``, then keeps that subprocess for the life of the + Agent API. ``POST /api/profile/switch`` is explicitly ``process_wide=False`` + — it sets a cookie / thread-local for sessions, but does **not** respawn + MCP with the target profile's ``DEEPSQL_AUTH_TOKEN``. + + The MCP client re-reads ``DEEPSQL_TOKEN_FILE`` per request (mtime cache). + Writing only ``profiles/u-/deepsql.token`` therefore leaves the + already-running server on the previous user's credential (in practice the + admin who first opened the Agent tab). View as / a new chat thread then + executes SQL as admin and skips chat-access policy. + + Mirror onto ``$HERMES_HOME/deepsql.token`` and every profile token file so + whichever path the live process was started with picks up the rotation. + """ + written = [write_token_file(hermes_home, token)] + for home in iter_profile_homes(hermes_home): + written.append(write_token_file(home, token)) + return written + + def ensure_profile(name: str) -> Path: home = HERMES_HOME / "profiles" / name if home.exists(): @@ -210,6 +245,11 @@ def _handle_provision(self): token_file = write_token_file(home, token) write_profile_mcp(home, user=user, token=token, token_file=token_file) write_profile_env(home, user=user, token=token, token_file=token_file) + mirrored = mirror_token_for_live_mcp(HERMES_HOME, token) + sys.stderr.write( + f"[agent-provisioner] mirrored MCP token for {user} onto " + f"{len(mirrored)} live token file(s)\n" + ) except Exception as e: return self._send(500, {"error": str(e)}) return self._send(200, {"ok": True, "profile": profile, "home": str(home)}) @@ -243,6 +283,44 @@ def _handle_revoke(self): return self._send(200, {"ok": True, "profile": profile, "home": str(home)}) +def _self_test() -> int: + """Prove token mirroring covers the process-global MCP watch path.""" + import tempfile + import traceback + + hermes_home = Path(tempfile.mkdtemp(prefix="deepsql-provisioner-test-")) + try: + admin = hermes_home / "profiles" / "u-admin" + editor = hermes_home / "profiles" / "u-marts-editor" + admin.mkdir(parents=True) + editor.mkdir(parents=True) + write_token_file(admin, "dsql_mcp_admin.old") + write_token_file(editor, "dsql_mcp_editor.old") + + rotated = "dsql_mcp_editor.live" + written = mirror_token_for_live_mcp(hermes_home, rotated) + paths = {p.resolve() for p in written} + expected = { + token_file_for(hermes_home).resolve(), + token_file_for(admin).resolve(), + token_file_for(editor).resolve(), + } + if paths != expected: + raise AssertionError(f"mirrored paths {paths} != {expected}") + for path in expected: + body = path.read_text().strip() + if body != rotated: + raise AssertionError(f"{path} still {body!r}, expected {rotated!r}") + print("ok: live MCP token is mirrored onto every profile token file") + return 0 + except Exception: + traceback.print_exc() + return 1 + finally: + import shutil + shutil.rmtree(hermes_home, ignore_errors=True) + + def main(): if not SECRET: print("AGENT_PROVISION_SECRET is required", file=sys.stderr) @@ -258,4 +336,6 @@ def main(): if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--self-test": + sys.exit(_self_test()) main() diff --git a/src/components/AgentChat/AgentChatPanel.jsx b/src/components/AgentChat/AgentChatPanel.jsx index 3f80130..0326c6b 100644 --- a/src/components/AgentChat/AgentChatPanel.jsx +++ b/src/components/AgentChat/AgentChatPanel.jsx @@ -2,6 +2,7 @@ import { useState, useRef, useEffect, useCallback } from 'react' import { ArrowUp, Plus, Square, Loader2, Database, Sparkles, Hash, Table2, Clock, AlertCircle } from 'lucide-react' import { agentChatAPI, withConnectionContext } from '@/lib/api/agentClient' import { agentConversationAPI } from '@/lib/api/client' +import { useAuth } from '@/hooks/useAuth' import AgentMarkdown from './AgentMarkdown' import { sanitizeAssistantAnswer } from './sanitizeAssistantAnswer' import styles from './AgentChatPanel.module.css' @@ -41,6 +42,7 @@ const SUGGESTIONS = [ ] export default function AgentChatPanel({ connectionId, connectionName }) { + const { username } = useAuth() const [sessionId, setSessionId] = useState(null) const [booting, setBooting] = useState(true) const [bootError, setBootError] = useState(null) @@ -113,7 +115,7 @@ export default function AgentChatPanel({ connectionId, connectionName }) { } finally { setBooting(false) } - }, [connectionId]) + }, [connectionId, username]) useEffect(() => { boot() }, [boot]) useEffect(() => { if (listRef.current) listRef.current.scrollTop = listRef.current.scrollHeight }, [messages]) diff --git a/src/components/sections/AgentChatSection.jsx b/src/components/sections/AgentChatSection.jsx index eb156e0..9ae20fc 100644 --- a/src/components/sections/AgentChatSection.jsx +++ b/src/components/sections/AgentChatSection.jsx @@ -1,8 +1,10 @@ import { useConnectionManager } from '@/lib/hooks/useConnectionManager' +import { useAuth } from '@/hooks/useAuth' import AgentChatPanel from '@/components/AgentChat/AgentChatPanel' export default function AgentChatSection() { const { connectionId, selectedConnection } = useConnectionManager() + const { username } = useAuth() if (!connectionId) { return ( @@ -12,10 +14,11 @@ export default function AgentChatSection() { ) } - // Remount on connection change so the chat re-bootstraps a fresh session. + // Remount on connection *or* identity change so View as re-bootstraps the + // target user's agent profile instead of keeping the admin MCP session. return ( diff --git a/src/components/sections/ShareMenu.jsx b/src/components/sections/ShareMenu.jsx index cd343ea..fa3ae1f 100644 --- a/src/components/sections/ShareMenu.jsx +++ b/src/components/sections/ShareMenu.jsx @@ -16,6 +16,7 @@ export default function ShareMenu({ savedId, initialPublic, initialToken, initia const [pwInput, setPwInput] = useState('') const [pwEditing, setPwEditing] = useState(false) const [pwBusy, setPwBusy] = useState(false) + const [shareError, setShareError] = useState('') const ref = useRef(null) useEffect(() => { @@ -41,6 +42,7 @@ export default function ShareMenu({ savedId, initialPublic, initialToken, initia const togglePublic = async () => { if (busy || !savedId) return setBusy(true) + setShareError('') try { if (!pub) { const res = await savedDashboardsAPI.enableShare(savedId) @@ -48,7 +50,10 @@ export default function ShareMenu({ savedId, initialPublic, initialToken, initia } else { await savedDashboardsAPI.disableShare(savedId); setPub(false); onPublicChange?.(false) } - } catch { /* leave state as-is on failure */ } finally { setBusy(false) } + } catch (err) { + const message = err?.response?.data?.message || err?.message || 'Failed to update sharing' + setShareError(message) + } finally { setBusy(false) } } const savePassword = async () => { @@ -105,6 +110,7 @@ export default function ShareMenu({ savedId, initialPublic, initialToken, initia {busy && } + {shareError ?
{shareError}
: null} {pub ? ( <>
Anyone with this link can view — read-only, no login. Turn off to revoke.
diff --git a/src/components/sections/ShareMenu.module.css b/src/components/sections/ShareMenu.module.css index 3c1ee8b..29ea279 100644 --- a/src/components/sections/ShareMenu.module.css +++ b/src/components/sections/ShareMenu.module.css @@ -42,6 +42,7 @@ color: #111318; } .hint { font-size: 12px; color: #8b909b; line-height: 1.4; } +.error { font-size: 12px; color: #b91c1c; line-height: 1.4; } .linkRow { display: flex; gap: 6px; align-items: center; } .input { diff --git a/src/components/tabs/admin/UsersTab.jsx b/src/components/tabs/admin/UsersTab.jsx index ff9a4d5..1ab4a5a 100644 --- a/src/components/tabs/admin/UsersTab.jsx +++ b/src/components/tabs/admin/UsersTab.jsx @@ -1086,10 +1086,17 @@ function ConnectionAccessModal({ user, data, loading, onClose, onSave, onRevoke, Assign the connection first, then save a per-user chat policy for it. )} + {assignment?.chatAccessPolicy?.allowedSchemas?.length > 0 && ( +
+ Resolved scope:{' '} + {assignment.chatAccessPolicy.allowedSchemas.join(', ')} +
+ )} {preview && (
Blocked sensitivity: {(preview.blockedSensitivityCategories || []).join(', ') || 'None'}
+
Allowed schemas: {(preview.allowedSchemas || []).join(', ') || 'All schemas'}
Impacted tables: {(preview.impactedTables || []).join(', ') || 'None'}
Impacted columns: {(preview.impactedColumns || []).join(', ') || 'None'}
Modes: {preview.blockMode ? 'Block' : 'Allow'} + {preview.redactMode ? 'Redact' : 'Pass through'}
diff --git a/src/hooks/useAuth.jsx b/src/hooks/useAuth.jsx index 1a02762..2bd1663 100644 --- a/src/hooks/useAuth.jsx +++ b/src/hooks/useAuth.jsx @@ -2,6 +2,7 @@ import { useState, useEffect, createContext, useContext, useMemo, useCallback } import { useNavigate, useLocation } from 'react-router-dom' import { getActionPermission, getActionConfig } from '@/lib/actions' import { authAPI, setupAPI, adminAPI, AUTH_CHANGE_EVENT } from '@/lib/api/client' +import { clearAgentRemoteUser } from '@/lib/api/agentClient' import { queryClient } from '@/lib/queryClient' import { useChatStore } from '@/lib/stores/useChatStore' import { useConnectionStore } from '@/lib/stores/useConnectionStore' @@ -205,12 +206,14 @@ export function AuthProvider({ children }) { const startImpersonation = useCallback(async (userId) => { const payload = await adminAPI.startImpersonation(userId) + clearAgentRemoteUser() applyAuthPayload(payload, { resetSession: true }) return payload }, [applyAuthPayload]) const stopImpersonation = useCallback(async () => { const payload = await adminAPI.stopImpersonation() + clearAgentRemoteUser() applyAuthPayload(payload, { resetSession: true }) return payload }, [applyAuthPayload]) diff --git a/src/lib/api/agentClient.js b/src/lib/api/agentClient.js index cde1eca..65853cb 100644 --- a/src/lib/api/agentClient.js +++ b/src/lib/api/agentClient.js @@ -27,6 +27,11 @@ let agentCsrfToken = null; */ let agentRemoteUser = null; +export function clearAgentRemoteUser() { + agentRemoteUser = null + agentCsrfToken = null +} + function withAgentAuthHeaders(headers = {}) { if (agentRemoteUser) { headers["X-Remote-User"] = agentRemoteUser; @@ -165,8 +170,10 @@ export const agentChatAPI = { /** Subscribe to a turn's SSE stream. Returns the EventSource (caller may .close()). */ streamChat(streamId, { onToken, onTool, onToolComplete, onEnd, onError } = {}) { + const qs = new URLSearchParams({ stream_id: streamId }) + if (agentRemoteUser) qs.set('remote_user', agentRemoteUser) const es = new EventSource( - `${AGENT_BASE}/api/chat/stream?stream_id=${encodeURIComponent(streamId)}`, + `${AGENT_BASE}/api/chat/stream?${qs.toString()}`, { withCredentials: true }, ); let done = false; diff --git a/vite.config.js b/vite.config.js index e20bbbf..154fffb 100644 --- a/vite.config.js +++ b/vite.config.js @@ -73,6 +73,17 @@ export default defineConfig({ if (typeof incoming === 'string' && incoming.trim()) { lastRemoteUser = incoming.trim() } + try { + const url = new URL(req.url, 'http://vite.local') + const fromQuery = url.searchParams.get('remote_user') + if (fromQuery && fromQuery.trim()) { + lastRemoteUser = fromQuery.trim() + url.searchParams.delete('remote_user') + proxyReq.path = url.pathname + url.search + } + } catch { + /* keep lastRemoteUser */ + } if (lastRemoteUser) { proxyReq.setHeader('X-Remote-User', lastRemoteUser) }