pmcp_code_mode/
validation.rs

1//! Validation pipeline for Code Mode.
2//!
3//! The pipeline validates code through multiple stages:
4//! 1. Parse (syntax check)
5//! 2. Policy evaluation (PolicyEvaluator trait or basic config checks)
6//! 3. Security analysis
7//! 4. Explanation generation
8//! 5. Token generation
9
10use crate::config::CodeModeConfig;
11#[cfg(feature = "openapi-code-mode")]
12use crate::config::OperationRegistry;
13use crate::explanation::{ExplanationGenerator, TemplateExplanationGenerator};
14use crate::graphql::{GraphQLQueryInfo, GraphQLValidator};
15use crate::policy::{OperationEntity, PolicyEvaluator};
16use crate::token::{compute_context_hash, HmacTokenGenerator, TokenGenerator, TokenSecret};
17use crate::types::{
18    PolicyViolation, TokenError, UnifiedAction, ValidationError, ValidationMetadata,
19    ValidationResult,
20};
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::sync::Arc;
23use std::time::Instant;
24
25#[cfg(feature = "openapi-code-mode")]
26use crate::javascript::{JavaScriptCodeInfo, JavaScriptValidator};
27
28/// Static flag to ensure the "no policy evaluator" warning is only logged once per process.
29static NO_POLICY_WARNING_LOGGED: AtomicBool = AtomicBool::new(false);
30
31/// Log a warning when Code Mode is enabled without a policy evaluator.
32fn warn_no_policy_configured() {
33    if !NO_POLICY_WARNING_LOGGED.swap(true, Ordering::SeqCst) {
34        tracing::warn!(
35            target: "code_mode",
36            "CODE MODE SECURITY WARNING: Code Mode is enabled but no policy evaluator \
37            is configured. Only basic config checks (allow_mutations, max_depth, etc.) will be \
38            performed. This provides NO real authorization policy evaluation. \
39            For production deployments, configure a policy evaluator (AVP or local Cedar)."
40        );
41    }
42}
43
44/// Context for validation (user, session, schema).
45#[derive(Debug, Clone)]
46pub struct ValidationContext {
47    /// User ID from access token
48    pub user_id: String,
49
50    /// MCP session ID
51    pub session_id: String,
52
53    /// Schema hash for context binding
54    pub schema_hash: String,
55
56    /// Permissions hash for context binding
57    pub permissions_hash: String,
58}
59
60impl ValidationContext {
61    /// Create a new validation context.
62    pub fn new(
63        user_id: impl Into<String>,
64        session_id: impl Into<String>,
65        schema_hash: impl Into<String>,
66        permissions_hash: impl Into<String>,
67    ) -> Self {
68        Self {
69            user_id: user_id.into(),
70            session_id: session_id.into(),
71            schema_hash: schema_hash.into(),
72            permissions_hash: permissions_hash.into(),
73        }
74    }
75
76    /// Compute the combined context hash.
77    pub fn context_hash(&self) -> String {
78        compute_context_hash(&self.schema_hash, &self.permissions_hash)
79    }
80}
81
82/// The validation pipeline that orchestrates all validation stages.
83pub struct ValidationPipeline<
84    T: TokenGenerator = HmacTokenGenerator,
85    E: ExplanationGenerator = TemplateExplanationGenerator,
86> {
87    config: CodeModeConfig,
88    graphql_validator: GraphQLValidator,
89    #[cfg(feature = "openapi-code-mode")]
90    javascript_validator: JavaScriptValidator,
91    #[cfg(feature = "openapi-code-mode")]
92    operation_registry: OperationRegistry,
93    token_generator: T,
94    explanation_generator: E,
95    policy_evaluator: Option<Arc<dyn PolicyEvaluator>>,
96}
97
98impl ValidationPipeline<HmacTokenGenerator, TemplateExplanationGenerator> {
99    /// Create a new validation pipeline with default generators.
100    ///
101    /// **Warning**: This constructor does not configure a policy evaluator.
102    /// Only basic config checks will be performed.
103    ///
104    /// # Errors
105    ///
106    /// Returns [`TokenError::SecretTooShort`] if the token secret is shorter
107    /// than [`HmacTokenGenerator::MIN_SECRET_LEN`] (16 bytes).
108    pub fn new(
109        config: CodeModeConfig,
110        token_secret: impl Into<Vec<u8>>,
111    ) -> Result<Self, TokenError> {
112        if config.enabled {
113            warn_no_policy_configured();
114        }
115
116        #[cfg(feature = "openapi-code-mode")]
117        let operation_registry = OperationRegistry::from_entries(&config.operations);
118
119        Ok(Self {
120            graphql_validator: GraphQLValidator::default(),
121            #[cfg(feature = "openapi-code-mode")]
122            javascript_validator: JavaScriptValidator::default()
123                .with_sdk_operations(config.sdk_operations.clone()),
124            #[cfg(feature = "openapi-code-mode")]
125            operation_registry,
126            token_generator: HmacTokenGenerator::new_from_bytes(token_secret)?,
127            explanation_generator: TemplateExplanationGenerator::new(),
128            policy_evaluator: None,
129            config,
130        })
131    }
132
133    /// Create a new validation pipeline from a [`TokenSecret`].
134    ///
135    /// Convenience constructor for production callers and derive macro generated
136    /// code. Callers never need to call `expose_secret()` directly.
137    ///
138    /// **Security note**: Internally this creates an intermediate `Vec<u8>` copy
139    /// of the secret bytes that is **not** zeroized on drop. For maximum security,
140    /// prefer [`TokenSecret::from_env`] which minimizes secret copies. This
141    /// limitation will be addressed in a future version by adding a
142    /// `HmacTokenGenerator::from_secret_ref` constructor.
143    ///
144    /// **Warning**: This constructor does not configure a policy evaluator.
145    /// Only basic config checks will be performed.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`TokenError::SecretTooShort`] if the token secret is shorter
150    /// than [`HmacTokenGenerator::MIN_SECRET_LEN`] (16 bytes).
151    pub fn from_token_secret(
152        config: CodeModeConfig,
153        secret: &TokenSecret,
154    ) -> Result<Self, TokenError> {
155        Self::new(config, secret.expose_secret().to_vec())
156    }
157
158    /// Create a new validation pipeline with a policy evaluator.
159    ///
160    /// # Errors
161    ///
162    /// Returns [`TokenError::SecretTooShort`] if the token secret is shorter
163    /// than [`HmacTokenGenerator::MIN_SECRET_LEN`] (16 bytes).
164    pub fn with_policy_evaluator(
165        config: CodeModeConfig,
166        token_secret: impl Into<Vec<u8>>,
167        evaluator: Arc<dyn PolicyEvaluator>,
168    ) -> Result<Self, TokenError> {
169        #[cfg(feature = "openapi-code-mode")]
170        let operation_registry = OperationRegistry::from_entries(&config.operations);
171
172        Ok(Self {
173            graphql_validator: GraphQLValidator::default(),
174            #[cfg(feature = "openapi-code-mode")]
175            javascript_validator: JavaScriptValidator::default()
176                .with_sdk_operations(config.sdk_operations.clone()),
177            #[cfg(feature = "openapi-code-mode")]
178            operation_registry,
179            token_generator: HmacTokenGenerator::new_from_bytes(token_secret)?,
180            explanation_generator: TemplateExplanationGenerator::new(),
181            policy_evaluator: Some(evaluator),
182            config,
183        })
184    }
185
186    /// Create a pipeline from a [`TokenSecret`] with an `Arc` policy evaluator.
187    ///
188    /// Used by derive macro generated code where the policy evaluator is
189    /// stored as `Arc<dyn PolicyEvaluator>` on the parent struct.
190    ///
191    /// # Errors
192    ///
193    /// Returns [`TokenError::SecretTooShort`] if the token secret is shorter
194    /// than [`HmacTokenGenerator::MIN_SECRET_LEN`] (16 bytes).
195    pub fn from_token_secret_with_policy(
196        config: CodeModeConfig,
197        secret: &TokenSecret,
198        evaluator: Arc<dyn PolicyEvaluator>,
199    ) -> Result<Self, TokenError> {
200        Self::with_policy_evaluator(config, secret.expose_secret().to_vec(), evaluator)
201    }
202}
203
204impl<T: TokenGenerator, E: ExplanationGenerator> ValidationPipeline<T, E> {
205    /// Create a pipeline with custom generators.
206    pub fn with_generators(
207        config: CodeModeConfig,
208        token_generator: T,
209        explanation_generator: E,
210    ) -> Self {
211        #[cfg(feature = "openapi-code-mode")]
212        let operation_registry = OperationRegistry::from_entries(&config.operations);
213
214        Self {
215            graphql_validator: GraphQLValidator::default(),
216            #[cfg(feature = "openapi-code-mode")]
217            javascript_validator: JavaScriptValidator::default()
218                .with_sdk_operations(config.sdk_operations.clone()),
219            #[cfg(feature = "openapi-code-mode")]
220            operation_registry,
221            token_generator,
222            explanation_generator,
223            policy_evaluator: None,
224            config,
225        }
226    }
227
228    /// Set the policy evaluator for this pipeline.
229    pub fn set_policy_evaluator(&mut self, evaluator: Arc<dyn PolicyEvaluator>) {
230        self.policy_evaluator = Some(evaluator);
231    }
232
233    /// Check if a policy evaluator is configured.
234    pub fn has_policy_evaluator(&self) -> bool {
235        self.policy_evaluator.is_some()
236    }
237
238    /// Check mutation and query authorization against config (blocklists, allowlists).
239    ///
240    /// This is the authorization logic shared between the sync and async validation paths.
241    /// It uses the already-parsed `query_info` to avoid re-parsing.
242    fn check_config_authorization(
243        &self,
244        query_info: &GraphQLQueryInfo,
245        start: Instant,
246    ) -> Option<ValidationResult> {
247        // Mutation authorization checks
248        if !query_info.operation_type.is_read_only() {
249            let mutation_name = query_info.root_fields.first().cloned().unwrap_or_default();
250
251            if !self.config.blocked_mutations.is_empty()
252                && self.config.blocked_mutations.contains(&mutation_name)
253            {
254                return Some(ValidationResult::failure(
255                    vec![PolicyViolation::new(
256                        "code_mode",
257                        "blocked_mutation",
258                        &format!("Mutation '{}' is blocked for this server", mutation_name),
259                    )
260                    .with_suggestion("This mutation is in the blocklist and cannot be executed")],
261                    self.build_metadata(query_info, start.elapsed().as_millis() as u64),
262                ));
263            }
264
265            if !self.config.allowed_mutations.is_empty() {
266                if !self.config.allowed_mutations.contains(&mutation_name) {
267                    return Some(ValidationResult::failure(
268                        vec![PolicyViolation::new(
269                            "code_mode",
270                            "mutation_not_allowed",
271                            &format!("Mutation '{}' is not in the allowlist", mutation_name),
272                        )
273                        .with_suggestion(&format!(
274                            "Only these mutations are allowed: {}",
275                            self.config
276                                .allowed_mutations
277                                .iter()
278                                .cloned()
279                                .collect::<Vec<_>>()
280                                .join(", ")
281                        ))],
282                        self.build_metadata(query_info, start.elapsed().as_millis() as u64),
283                    ));
284                }
285            } else if !self.config.allow_mutations {
286                return Some(ValidationResult::failure(
287                    vec![PolicyViolation::new(
288                        "code_mode",
289                        "allow_mutations",
290                        "Mutations are not enabled for this server",
291                    )
292                    .with_suggestion("Only read-only queries are allowed")],
293                    self.build_metadata(query_info, start.elapsed().as_millis() as u64),
294                ));
295            }
296        }
297
298        // Query (read) authorization checks -- mirrors mutation enforcement above
299        if query_info.operation_type.is_read_only() {
300            let query_name = query_info.root_fields.first().cloned().unwrap_or_default();
301
302            if !self.config.blocked_queries.is_empty()
303                && self.config.blocked_queries.contains(&query_name)
304            {
305                return Some(ValidationResult::failure(
306                    vec![PolicyViolation::new(
307                        "code_mode",
308                        "blocked_query",
309                        &format!("Query '{}' is blocked for this server", query_name),
310                    )
311                    .with_suggestion("This query is in the blocklist and cannot be executed")],
312                    self.build_metadata(query_info, start.elapsed().as_millis() as u64),
313                ));
314            }
315
316            if !self.config.allowed_queries.is_empty()
317                && !self.config.allowed_queries.contains(&query_name)
318            {
319                return Some(ValidationResult::failure(
320                    vec![PolicyViolation::new(
321                        "code_mode",
322                        "query_not_allowed",
323                        &format!("Query '{}' is not in the allowlist", query_name),
324                    )
325                    .with_suggestion(&format!(
326                        "Only these queries are allowed: {}",
327                        self.config
328                            .allowed_queries
329                            .iter()
330                            .cloned()
331                            .collect::<Vec<_>>()
332                            .join(", ")
333                    ))],
334                    self.build_metadata(query_info, start.elapsed().as_millis() as u64),
335                ));
336            }
337        }
338
339        None
340    }
341
342    /// Validate a GraphQL query using basic config checks only.
343    pub fn validate_graphql_query(
344        &self,
345        query: &str,
346        context: &ValidationContext,
347    ) -> Result<ValidationResult, ValidationError> {
348        let start = Instant::now();
349
350        if !self.config.enabled {
351            return Err(ValidationError::ConfigError(
352                "Code Mode is not enabled for this server".into(),
353            ));
354        }
355
356        if query.len() > self.config.max_query_length {
357            return Err(ValidationError::SecurityError {
358                message: format!(
359                    "Query length {} exceeds maximum {}",
360                    query.len(),
361                    self.config.max_query_length
362                ),
363                issue: crate::types::SecurityIssueType::HighComplexity,
364            });
365        }
366
367        let query_info = self.graphql_validator.validate(query)?;
368
369        // Config-based authorization checks (mutation blocklist/allowlist, query blocklist/allowlist)
370        if let Some(failure) = self.check_config_authorization(&query_info, start) {
371            return Ok(failure);
372        }
373
374        self.complete_validation(query, &query_info, context, start)
375    }
376
377    /// Validate a GraphQL query using a policy evaluator (async).
378    pub async fn validate_graphql_query_async(
379        &self,
380        query: &str,
381        context: &ValidationContext,
382    ) -> Result<ValidationResult, ValidationError> {
383        let start = Instant::now();
384
385        if !self.config.enabled {
386            return Err(ValidationError::ConfigError(
387                "Code Mode is not enabled for this server".into(),
388            ));
389        }
390
391        if query.len() > self.config.max_query_length {
392            return Err(ValidationError::SecurityError {
393                message: format!(
394                    "Query length {} exceeds maximum {}",
395                    query.len(),
396                    self.config.max_query_length
397                ),
398                issue: crate::types::SecurityIssueType::HighComplexity,
399            });
400        }
401
402        let query_info = self.graphql_validator.validate(query)?;
403
404        // Policy evaluation via trait
405        if let Some(ref evaluator) = self.policy_evaluator {
406            let operation_entity = OperationEntity::from_query_info(&query_info);
407            let server_config = self.config.to_server_config_entity();
408
409            let decision = evaluator
410                .evaluate_operation(&operation_entity, &server_config)
411                .await
412                .map_err(|e| {
413                    ValidationError::InternalError(format!("Policy evaluation error: {}", e))
414                })?;
415
416            if !decision.allowed {
417                let violations: Vec<PolicyViolation> = decision
418                    .determining_policies
419                    .iter()
420                    .map(|policy_id| {
421                        PolicyViolation::new(
422                            "policy",
423                            policy_id.clone(),
424                            "Policy denied the operation",
425                        )
426                    })
427                    .collect();
428
429                return Ok(ValidationResult::failure(
430                    violations,
431                    self.build_metadata(&query_info, start.elapsed().as_millis() as u64),
432                ));
433            }
434        } else {
435            warn_no_policy_configured();
436            tracing::debug!(
437                target: "code_mode",
438                "Falling back to basic config checks (no policy evaluator configured)"
439            );
440            // Reuse already-parsed query_info instead of re-parsing via validate_graphql_query
441            if let Some(failure) = self.check_config_authorization(&query_info, start) {
442                return Ok(failure);
443            }
444        }
445
446        self.complete_validation(query, &query_info, context, start)
447    }
448
449    /// Complete validation after policy check passes.
450    fn complete_validation(
451        &self,
452        query: &str,
453        query_info: &GraphQLQueryInfo,
454        context: &ValidationContext,
455        start: Instant,
456    ) -> Result<ValidationResult, ValidationError> {
457        let security_analysis = self.graphql_validator.analyze_security(query_info);
458        let risk_level = security_analysis.assess_risk();
459
460        if security_analysis
461            .potential_issues
462            .iter()
463            .any(|i| i.is_critical())
464        {
465            let violations: Vec<PolicyViolation> = security_analysis
466                .potential_issues
467                .iter()
468                .filter(|i| i.is_critical())
469                .map(|i| {
470                    PolicyViolation::new("security", format!("{:?}", i.issue_type), &i.message)
471                })
472                .collect();
473
474            return Ok(ValidationResult::failure(
475                violations,
476                self.build_metadata(query_info, start.elapsed().as_millis() as u64),
477            ));
478        }
479
480        let explanation = self
481            .explanation_generator
482            .explain_graphql(query_info, &security_analysis);
483
484        let context_hash = context.context_hash();
485        let token = self.token_generator.generate(
486            query,
487            &context.user_id,
488            &context.session_id,
489            self.config.server_id(),
490            &context_hash,
491            risk_level,
492            self.config.token_ttl_seconds,
493        );
494
495        let token_string = token.encode().map_err(|e| {
496            ValidationError::InternalError(format!("Failed to encode token: {}", e))
497        })?;
498
499        let operation_type_str = format!("{:?}", query_info.operation_type).to_lowercase();
500        let mutation_name = query_info.operation_name.as_deref();
501        let inferred_action = UnifiedAction::from_graphql(&operation_type_str, mutation_name);
502        let action = UnifiedAction::resolve(
503            inferred_action,
504            &self.config.action_tags,
505            query_info.operation_name.as_deref().unwrap_or(""),
506        );
507
508        let metadata = ValidationMetadata {
509            is_read_only: query_info.operation_type.is_read_only(),
510            estimated_rows: security_analysis.estimated_rows,
511            accessed_types: security_analysis.tables_accessed.iter().cloned().collect(),
512            accessed_fields: security_analysis.fields_accessed.iter().cloned().collect(),
513            has_aggregation: security_analysis.has_aggregation,
514            code_type: Some(self.graphql_validator.to_code_type(query_info)),
515            action: Some(action),
516            validation_time_ms: start.elapsed().as_millis() as u64,
517        };
518
519        let mut result = ValidationResult::success(explanation, risk_level, token_string, metadata);
520
521        for issue in &security_analysis.potential_issues {
522            if !issue.is_critical() {
523                result.warnings.push(issue.message.clone());
524            }
525        }
526
527        Ok(result)
528    }
529
530    /// Build metadata from query info.
531    fn build_metadata(
532        &self,
533        query_info: &GraphQLQueryInfo,
534        validation_time_ms: u64,
535    ) -> ValidationMetadata {
536        let operation_type_str = format!("{:?}", query_info.operation_type).to_lowercase();
537        let mutation_name = query_info.operation_name.as_deref();
538        let inferred_action = UnifiedAction::from_graphql(&operation_type_str, mutation_name);
539        let action = UnifiedAction::resolve(
540            inferred_action,
541            &self.config.action_tags,
542            query_info.operation_name.as_deref().unwrap_or(""),
543        );
544
545        ValidationMetadata {
546            is_read_only: query_info.operation_type.is_read_only(),
547            estimated_rows: None,
548            accessed_types: query_info.types_accessed.iter().cloned().collect(),
549            accessed_fields: query_info.fields_accessed.iter().cloned().collect(),
550            has_aggregation: false,
551            code_type: Some(self.graphql_validator.to_code_type(query_info)),
552            action: Some(action),
553            validation_time_ms,
554        }
555    }
556
557    /// Validate JavaScript code for OpenAPI Code Mode (sync, no policy evaluation).
558    ///
559    /// Runs config-level checks only. For policy evaluation (Cedar/AVP), use
560    /// [`validate_javascript_code_async`] instead. Retained for backward
561    /// compatibility with callers that don't need policy enforcement.
562    #[cfg(feature = "openapi-code-mode")]
563    pub fn validate_javascript_code(
564        &self,
565        code: &str,
566        context: &ValidationContext,
567    ) -> Result<ValidationResult, ValidationError> {
568        let start = Instant::now();
569        let code_info = self.validate_js_preamble(code)?;
570        if let Some(failure) = self.check_js_config_authorization(&code_info, start) {
571            return Ok(failure);
572        }
573        self.complete_js_validation(code, &code_info, context, start)
574    }
575
576    /// Validate JavaScript code with async policy evaluation.
577    ///
578    /// Mirrors [`validate_graphql_query_async`] but for JavaScript/OpenAPI:
579    /// 1. Parse JS via SWC + config-level checks (shared with sync version)
580    /// 2. Policy evaluation via [`PolicyEvaluator::evaluate_script`] (async, fail-closed)
581    /// 3. Security analysis + token generation
582    ///
583    /// When no policy evaluator is configured, falls back to config-only checks.
584    #[cfg(feature = "openapi-code-mode")]
585    pub async fn validate_javascript_code_async(
586        &self,
587        code: &str,
588        context: &ValidationContext,
589    ) -> Result<ValidationResult, ValidationError> {
590        use crate::policy::types::ScriptEntity;
591
592        let start = Instant::now();
593        let code_info = self.validate_js_preamble(code)?;
594        if let Some(failure) = self.check_js_config_authorization(&code_info, start) {
595            return Ok(failure);
596        }
597
598        // Policy evaluation via evaluate_script (mirrors GraphQL's evaluate_operation)
599        if let Some(ref evaluator) = self.policy_evaluator {
600            let sensitive_patterns: Vec<String> =
601                self.config.openapi_blocked_paths.iter().cloned().collect();
602            let registry_ref = if self.operation_registry.is_empty() {
603                None
604            } else {
605                Some(&self.operation_registry)
606            };
607            let script_entity =
608                ScriptEntity::from_javascript_info(&code_info, &sensitive_patterns, registry_ref);
609            let server_entity = self.config.to_openapi_server_entity();
610
611            let decision = evaluator
612                .evaluate_script(&script_entity, &server_entity)
613                .await
614                .map_err(|e| {
615                    ValidationError::InternalError(format!("Policy evaluation error: {}", e))
616                })?;
617
618            if !decision.allowed {
619                let violations: Vec<PolicyViolation> = decision
620                    .determining_policies
621                    .iter()
622                    .map(|policy_id| {
623                        PolicyViolation::new(
624                            "policy",
625                            policy_id.clone(),
626                            "Policy denied the script",
627                        )
628                    })
629                    .collect();
630
631                return Ok(ValidationResult::failure(
632                    violations,
633                    self.build_js_metadata(&code_info, start.elapsed().as_millis() as u64),
634                ));
635            }
636        }
637
638        self.complete_js_validation(code, &code_info, context, start)
639    }
640
641    /// Shared JavaScript preamble: enabled check, length check, parse.
642    #[cfg(feature = "openapi-code-mode")]
643    fn validate_js_preamble(&self, code: &str) -> Result<JavaScriptCodeInfo, ValidationError> {
644        if !self.config.enabled {
645            return Err(ValidationError::ConfigError(
646                "Code Mode is not enabled for this server".into(),
647            ));
648        }
649
650        if code.len() > self.config.max_query_length {
651            return Err(ValidationError::SecurityError {
652                message: format!(
653                    "Code length {} exceeds maximum {}",
654                    code.len(),
655                    self.config.max_query_length
656                ),
657                issue: crate::types::SecurityIssueType::HighComplexity,
658            });
659        }
660
661        Ok(self.javascript_validator.validate(code)?)
662    }
663
664    /// Config-level authorization checks for JavaScript code.
665    ///
666    /// Returns `Some(failure)` if a config check denied the code,
667    /// `None` if all checks pass. Shared between sync and async paths
668    /// (mirrors `check_config_authorization` for GraphQL).
669    #[cfg(feature = "openapi-code-mode")]
670    fn check_js_config_authorization(
671        &self,
672        code_info: &JavaScriptCodeInfo,
673        start: Instant,
674    ) -> Option<ValidationResult> {
675        if code_info.is_read_only {
676            return None;
677        }
678
679        for method in &code_info.methods_used {
680            if !self.config.openapi_blocked_writes.is_empty()
681                && self.config.openapi_blocked_writes.contains(method)
682            {
683                return Some(ValidationResult::failure(
684                    vec![PolicyViolation::new(
685                        "code_mode",
686                        "blocked_method",
687                        &format!("HTTP method '{}' is blocked for this server", method),
688                    )
689                    .with_suggestion("This method is in the blocklist and cannot be used")],
690                    self.build_js_metadata(code_info, start.elapsed().as_millis() as u64),
691                ));
692            }
693        }
694
695        if !self.config.openapi_allowed_writes.is_empty() {
696            tracing::debug!(
697                target: "code_mode",
698                "Skipping method-level check - policy evaluator will check operation allowlist ({} entries)",
699                self.config.openapi_allowed_writes.len()
700            );
701        } else if !self.config.openapi_allow_writes {
702            return Some(ValidationResult::failure(
703                vec![PolicyViolation::new(
704                    "code_mode",
705                    "allow_mutations",
706                    "Write HTTP methods (POST, PUT, DELETE, PATCH) are not enabled for this server",
707                )
708                .with_suggestion("Only read-only methods (GET, HEAD, OPTIONS) are allowed. Contact your administrator to enable write operations.")],
709                self.build_js_metadata(code_info, start.elapsed().as_millis() as u64),
710            ));
711        }
712
713        None
714    }
715
716    /// Complete JavaScript validation after policy checks pass.
717    #[cfg(feature = "openapi-code-mode")]
718    fn complete_js_validation(
719        &self,
720        code: &str,
721        code_info: &JavaScriptCodeInfo,
722        context: &ValidationContext,
723        start: Instant,
724    ) -> Result<ValidationResult, ValidationError> {
725        let security_analysis = self.javascript_validator.analyze_security(code_info);
726        let risk_level = security_analysis.assess_risk();
727
728        if security_analysis
729            .potential_issues
730            .iter()
731            .any(|i| i.is_critical())
732        {
733            let violations: Vec<PolicyViolation> = security_analysis
734                .potential_issues
735                .iter()
736                .filter(|i| i.is_critical())
737                .map(|i| {
738                    PolicyViolation::new("security", format!("{:?}", i.issue_type), &i.message)
739                })
740                .collect();
741
742            return Ok(ValidationResult::failure(
743                violations,
744                self.build_js_metadata(code_info, start.elapsed().as_millis() as u64),
745            ));
746        }
747
748        let explanation = self.generate_js_explanation(code_info, &security_analysis);
749
750        let context_hash = context.context_hash();
751        let token = self.token_generator.generate(
752            code,
753            &context.user_id,
754            &context.session_id,
755            self.config.server_id(),
756            &context_hash,
757            risk_level,
758            self.config.token_ttl_seconds,
759        );
760
761        let token_string = token.encode().map_err(|e| {
762            ValidationError::InternalError(format!("Failed to encode token: {}", e))
763        })?;
764
765        let metadata = self.build_js_metadata(code_info, start.elapsed().as_millis() as u64);
766
767        let mut result = ValidationResult::success(explanation, risk_level, token_string, metadata);
768
769        for issue in &security_analysis.potential_issues {
770            if !issue.is_critical() {
771                result.warnings.push(issue.message.clone());
772            }
773        }
774
775        Ok(result)
776    }
777
778    /// Build metadata from JavaScript code info.
779    #[cfg(feature = "openapi-code-mode")]
780    fn build_js_metadata(
781        &self,
782        code_info: &JavaScriptCodeInfo,
783        validation_time_ms: u64,
784    ) -> ValidationMetadata {
785        let action = if !code_info.api_calls.is_empty() {
786            let mut max_action = UnifiedAction::Read;
787            for call in &code_info.api_calls {
788                let method_str = format!("{:?}", call.method);
789                let inferred = UnifiedAction::from_http_method(&method_str);
790                match (&max_action, &inferred) {
791                    (UnifiedAction::Read, _) => max_action = inferred,
792                    (UnifiedAction::Write, UnifiedAction::Delete | UnifiedAction::Admin) => {
793                        max_action = inferred
794                    },
795                    (UnifiedAction::Delete, UnifiedAction::Admin) => max_action = inferred,
796                    _ => {},
797                }
798            }
799            Some(max_action)
800        } else if code_info.is_read_only {
801            Some(UnifiedAction::Read)
802        } else {
803            Some(UnifiedAction::Write)
804        };
805
806        ValidationMetadata {
807            is_read_only: code_info.is_read_only,
808            estimated_rows: None,
809            accessed_types: code_info.endpoints_accessed.iter().cloned().collect(),
810            accessed_fields: code_info.methods_used.iter().cloned().collect(),
811            has_aggregation: false,
812            code_type: Some(self.javascript_validator.to_code_type(code_info)),
813            action,
814            validation_time_ms,
815        }
816    }
817
818    /// Generate a human-readable explanation for JavaScript code.
819    #[cfg(feature = "openapi-code-mode")]
820    fn generate_js_explanation(
821        &self,
822        code_info: &JavaScriptCodeInfo,
823        security_analysis: &crate::types::SecurityAnalysis,
824    ) -> String {
825        let mut parts = Vec::new();
826
827        if code_info.is_read_only {
828            parts.push("This code will perform read-only API requests.".to_string());
829        } else {
830            parts.push("This code will perform API requests that may modify data.".to_string());
831        }
832
833        if !code_info.api_calls.is_empty() {
834            let call_descriptions: Vec<String> = code_info
835                .api_calls
836                .iter()
837                .map(|call| format!("{:?} {}", call.method, call.path))
838                .collect();
839
840            if call_descriptions.len() <= 3 {
841                parts.push(format!("API calls: {}", call_descriptions.join(", ")));
842            } else {
843                parts.push(format!(
844                    "API calls: {} and {} more",
845                    call_descriptions[..2].join(", "),
846                    call_descriptions.len() - 2
847                ));
848            }
849        }
850
851        if code_info.loop_count > 0 {
852            if code_info.all_loops_bounded {
853                parts.push(format!(
854                    "Contains {} bounded loop(s).",
855                    code_info.loop_count
856                ));
857            } else {
858                parts.push(format!(
859                    "Contains {} loop(s) - ensure they are properly bounded.",
860                    code_info.loop_count
861                ));
862            }
863        }
864
865        let risk = security_analysis.assess_risk();
866        parts.push(format!("Risk: {}", risk));
867
868        parts.join(" ")
869    }
870
871    /// Check if a validation result should be auto-approved.
872    pub fn should_auto_approve(&self, result: &ValidationResult) -> bool {
873        result.is_valid && self.config.should_auto_approve(result.risk_level)
874    }
875
876    /// Get the config.
877    pub fn config(&self) -> &CodeModeConfig {
878        &self.config
879    }
880
881    /// Get the token generator.
882    pub fn token_generator(&self) -> &T {
883        &self.token_generator
884    }
885}
886
887#[cfg(test)]
888mod tests {
889    use super::*;
890    use crate::types::RiskLevel;
891
892    fn test_pipeline() -> ValidationPipeline {
893        ValidationPipeline::new(CodeModeConfig::enabled(), b"test-secret-key!".to_vec()).unwrap()
894    }
895
896    fn test_context() -> ValidationContext {
897        ValidationContext::new("user-123", "session-456", "schema-hash", "perms-hash")
898    }
899
900    #[test]
901    fn test_simple_query_validation() {
902        let pipeline = test_pipeline();
903        let ctx = test_context();
904
905        let result = pipeline
906            .validate_graphql_query("query { users { id name } }", &ctx)
907            .unwrap();
908
909        assert!(result.is_valid);
910        assert!(result.approval_token.is_some());
911        assert_eq!(result.risk_level, RiskLevel::Low);
912        assert!(result.explanation.contains("read"));
913    }
914
915    #[test]
916    fn test_mutation_blocked() {
917        let mut config = CodeModeConfig::enabled();
918        config.allow_mutations = false;
919
920        let pipeline = ValidationPipeline::new(config, b"test-secret-key!".to_vec()).unwrap();
921        let ctx = test_context();
922
923        let result = pipeline
924            .validate_graphql_query("mutation { createUser(name: \"test\") { id } }", &ctx)
925            .unwrap();
926
927        assert!(!result.is_valid);
928        assert!(result
929            .violations
930            .iter()
931            .any(|v| v.rule == "allow_mutations"));
932    }
933
934    #[test]
935    fn test_disabled_code_mode() {
936        let config = CodeModeConfig::default();
937        let pipeline = ValidationPipeline::new(config, b"test-secret-key!".to_vec()).unwrap();
938        let ctx = test_context();
939
940        let result = pipeline.validate_graphql_query("query { users { id } }", &ctx);
941
942        assert!(matches!(result, Err(ValidationError::ConfigError(_))));
943    }
944
945    #[test]
946    fn test_auto_approve_low_risk() {
947        let pipeline = test_pipeline();
948        let ctx = test_context();
949
950        let result = pipeline
951            .validate_graphql_query("query { users { id } }", &ctx)
952            .unwrap();
953
954        assert!(pipeline.should_auto_approve(&result));
955    }
956
957    #[test]
958    fn test_context_hash() {
959        let ctx = test_context();
960        let hash1 = ctx.context_hash();
961
962        let ctx2 =
963            ValidationContext::new("user-123", "session-456", "different-schema", "perms-hash");
964        let hash2 = ctx2.context_hash();
965
966        assert_ne!(hash1, hash2);
967    }
968
969    #[test]
970    fn test_blocked_query_rejected() {
971        let mut config = CodeModeConfig::enabled();
972        config.blocked_queries.insert("users".to_string());
973
974        let pipeline = ValidationPipeline::new(config, b"test-secret-key!".to_vec()).unwrap();
975        let ctx = test_context();
976
977        let result = pipeline
978            .validate_graphql_query("query { users { id } }", &ctx)
979            .unwrap();
980
981        assert!(!result.is_valid);
982        assert!(result.violations.iter().any(|v| v.rule == "blocked_query"));
983    }
984
985    #[test]
986    fn test_allowed_queries_enforced() {
987        let mut config = CodeModeConfig::enabled();
988        config.allowed_queries.insert("orders".to_string());
989
990        let pipeline = ValidationPipeline::new(config, b"test-secret-key!".to_vec()).unwrap();
991        let ctx = test_context();
992
993        // "users" is not in the allowlist -- should be rejected
994        let result = pipeline
995            .validate_graphql_query("query { users { id } }", &ctx)
996            .unwrap();
997
998        assert!(!result.is_valid);
999        assert!(result
1000            .violations
1001            .iter()
1002            .any(|v| v.rule == "query_not_allowed"));
1003    }
1004}
pmcp_code_mode/validation.rs

pmcp_code_mode/
validation.rs
