1use std::net::SocketAddr;
20use std::sync::Arc;
21use std::sync::atomic::{AtomicU64, Ordering};
22use std::time::{Duration, Instant};
23
24use schemars::JsonSchema;
25use serde::Deserialize;
26use url::Url;
27
28use zeph_common::ToolName;
29
30use zeph_sanitizer::IpiFilter;
31
32use crate::audit::{AuditEntry, AuditLogger, AuditResult, EgressEvent, chrono_now};
33use crate::config::{EgressConfig, ScrapeConfig};
34use crate::executor::{
35 ClaimSource, ToolCall, ToolError, ToolExecutor, ToolOutput, deserialize_params,
36};
37use crate::net::{check_domain_policy, validate_url};
38
39fn redact_url_for_log(url: &str) -> String {
43 let Ok(mut parsed) = Url::parse(url) else {
44 return url.to_owned();
45 };
46 let _ = parsed.set_username("");
48 let _ = parsed.set_password(None);
49 let sensitive = [
51 "token", "key", "secret", "password", "auth", "sig", "api_key", "apikey",
52 ];
53 let filtered: Vec<(String, String)> = parsed
54 .query_pairs()
55 .filter(|(k, _)| {
56 let lower = k.to_lowercase();
57 !sensitive.iter().any(|s| lower.contains(s))
58 })
59 .map(|(k, v)| (k.into_owned(), v.into_owned()))
60 .collect();
61 if filtered.is_empty() {
62 parsed.set_query(None);
63 } else {
64 let q: String = filtered
65 .iter()
66 .map(|(k, v)| format!("{k}={v}"))
67 .collect::<Vec<_>>()
68 .join("&");
69 parsed.set_query(Some(&q));
70 }
71 parsed.to_string()
72}
73
74#[derive(Debug, Deserialize, JsonSchema)]
75struct FetchParams {
76 url: String,
78}
79
80#[derive(Debug, Deserialize, JsonSchema)]
81struct ScrapeInstruction {
82 url: String,
84 select: String,
86 #[serde(default = "default_extract")]
88 extract: String,
89 limit: Option<usize>,
91}
92
93fn default_extract() -> String {
94 "text".into()
95}
96
97#[derive(Debug)]
98enum ExtractMode {
99 Text,
100 Html,
101 Attr(String),
102}
103
104impl ExtractMode {
105 fn parse(s: &str) -> Self {
106 match s {
107 "text" => Self::Text,
108 "html" => Self::Html,
109 attr if attr.starts_with("attr:") => {
110 Self::Attr(attr.strip_prefix("attr:").unwrap_or(attr).to_owned())
111 }
112 _ => Self::Text,
113 }
114 }
115}
116
117#[derive(Debug)]
159pub struct WebScrapeExecutor {
160 timeout: Duration,
161 max_body_bytes: usize,
162 allowed_domains: Vec<String>,
163 denied_domains: Vec<String>,
164 audit_logger: Option<Arc<AuditLogger>>,
165 egress_config: EgressConfig,
166 egress_tx: Option<tokio::sync::mpsc::Sender<EgressEvent>>,
167 egress_dropped: Arc<AtomicU64>,
168 ipi_filter: IpiFilter,
170}
171
172impl WebScrapeExecutor {
173 #[must_use]
177 pub fn new(config: &ScrapeConfig) -> Self {
178 Self {
179 timeout: Duration::from_secs(config.timeout),
180 max_body_bytes: config.max_body_bytes,
181 allowed_domains: config.allowed_domains.clone(),
182 denied_domains: config.denied_domains.clone(),
183 audit_logger: None,
184 egress_config: EgressConfig::default(),
185 egress_tx: None,
186 egress_dropped: Arc::new(AtomicU64::new(0)),
187 ipi_filter: IpiFilter::new(config.ipi_filter_threshold),
188 }
189 }
190
191 #[must_use]
193 pub fn with_audit(mut self, logger: Arc<AuditLogger>) -> Self {
194 self.audit_logger = Some(logger);
195 self
196 }
197
198 #[must_use]
200 pub fn with_egress_config(mut self, config: EgressConfig) -> Self {
201 self.egress_config = config;
202 self
203 }
204
205 #[must_use]
210 pub fn with_egress_tx(
211 mut self,
212 tx: tokio::sync::mpsc::Sender<EgressEvent>,
213 dropped: Arc<AtomicU64>,
214 ) -> Self {
215 self.egress_tx = Some(tx);
216 self.egress_dropped = dropped;
217 self
218 }
219
220 #[must_use]
222 pub fn egress_dropped(&self) -> Arc<AtomicU64> {
223 Arc::clone(&self.egress_dropped)
224 }
225
226 fn build_client(&self, host: &str, addrs: &[SocketAddr]) -> reqwest::Client {
227 let mut builder = reqwest::Client::builder()
228 .timeout(self.timeout)
229 .redirect(reqwest::redirect::Policy::none());
230 builder = builder.resolve_to_addrs(host, addrs);
231 builder.build().unwrap_or_default()
232 }
233}
234
235impl ToolExecutor for WebScrapeExecutor {
236 fn tool_definitions(&self) -> Vec<crate::registry::ToolDef> {
237 use crate::registry::{InvocationHint, ToolDef};
238 vec![
239 ToolDef {
240 id: "web_scrape".into(),
241 description: "Extract structured data from a web page using CSS selectors.\n\nONLY call this tool when the user has explicitly provided a URL in their message, or when a prior tool call returned a URL to retrieve. NEVER construct, guess, or infer a URL from entity names, brand knowledge, or domain patterns.\n\nParameters: url (string, required) - HTTPS URL; select (string, required) - CSS selector; extract (string, optional) - \"text\", \"html\", or \"attr:<name>\"; limit (integer, optional) - max results\nReturns: extracted text/HTML/attribute values, one per line\nErrors: InvalidParams if URL is not HTTPS or selector is empty; Timeout after configured seconds; connection/DNS failures".into(),
242 schema: schemars::schema_for!(ScrapeInstruction),
243 invocation: InvocationHint::FencedBlock("scrape"),
244 output_schema: None,
245 server_id: None,
246 },
247 ToolDef {
248 id: "fetch".into(),
249 description: "Fetch a URL and return the response body as plain text.\n\nONLY call this tool when the user has explicitly provided a URL in their message, or when a prior tool call returned a URL to retrieve. NEVER construct, guess, or infer a URL from entity names, brand knowledge, or domain patterns. If no URL is present in the conversation, do not call this tool.\n\nParameters: url (string, required) - HTTPS URL to fetch\nReturns: response body as UTF-8 text, truncated if exceeding max body size\nErrors: InvalidParams if URL is not HTTPS; Timeout; SSRF-blocked private IPs; connection failures".into(),
250 schema: schemars::schema_for!(FetchParams),
251 invocation: InvocationHint::ToolCall,
252 output_schema: None,
253 server_id: None,
254 },
255 ]
256 }
257
258 async fn execute(&self, response: &str) -> Result<Option<ToolOutput>, ToolError> {
259 let blocks = extract_scrape_blocks(response);
260 if blocks.is_empty() {
261 return Ok(None);
262 }
263
264 let mut outputs = Vec::with_capacity(blocks.len());
265 #[allow(clippy::cast_possible_truncation)]
266 let blocks_executed = blocks.len() as u32;
267
268 for block in &blocks {
269 let instruction: ScrapeInstruction = serde_json::from_str(block).map_err(|e| {
270 ToolError::Execution(std::io::Error::new(
271 std::io::ErrorKind::InvalidData,
272 e.to_string(),
273 ))
274 })?;
275 let correlation_id = EgressEvent::new_correlation_id();
276 let start = Instant::now();
277 let scrape_result = self
278 .scrape_instruction(&instruction, &correlation_id, None, None)
279 .await;
280 #[allow(clippy::cast_possible_truncation)]
281 let duration_ms = start.elapsed().as_millis() as u64;
282 match scrape_result {
283 Ok(output) => {
284 self.log_audit(
285 "web_scrape",
286 &redact_url_for_log(&instruction.url),
287 AuditResult::Success,
288 duration_ms,
289 None,
290 None,
291 None,
292 Some(correlation_id),
293 )
294 .await;
295 outputs.push(output);
296 }
297 Err(e) => {
298 let audit_result = tool_error_to_audit_result(&e);
299 self.log_audit(
300 "web_scrape",
301 &redact_url_for_log(&instruction.url),
302 audit_result,
303 duration_ms,
304 Some(&e),
305 None,
306 None,
307 Some(correlation_id),
308 )
309 .await;
310 return Err(e);
311 }
312 }
313 }
314
315 Ok(Some(ToolOutput {
316 tool_name: ToolName::new("web-scrape"),
317 summary: outputs.join("\n\n"),
318 blocks_executed,
319 filter_stats: None,
320 diff: None,
321 streamed: false,
322 terminal_id: None,
323 locations: None,
324 raw_response: None,
325 claim_source: Some(ClaimSource::WebScrape),
326 ..Default::default()
327 }))
328 }
329
330 #[cfg_attr(
331 feature = "profiling",
332 tracing::instrument(name = "tools.scrape.fetch", skip_all)
333 )]
334 async fn execute_tool_call(&self, call: &ToolCall) -> Result<Option<ToolOutput>, ToolError> {
335 match call.tool_id.as_str() {
336 "web_scrape" => {
337 let instruction: ScrapeInstruction = deserialize_params(&call.params)?;
338 let correlation_id = EgressEvent::new_correlation_id();
339 let start = Instant::now();
340 let result = self
341 .scrape_instruction(
342 &instruction,
343 &correlation_id,
344 call.caller_id.clone(),
345 call.skill_name.clone(),
346 )
347 .await;
348 #[allow(clippy::cast_possible_truncation)]
349 let duration_ms = start.elapsed().as_millis() as u64;
350 self.run_with_audit(
351 "web_scrape",
352 "web-scrape",
353 &redact_url_for_log(&instruction.url),
354 call.caller_id.clone(),
355 call.skill_name.clone(),
356 correlation_id,
357 duration_ms,
358 result,
359 )
360 .await
361 }
362 "fetch" => {
363 let p: FetchParams = deserialize_params(&call.params)?;
364 let correlation_id = EgressEvent::new_correlation_id();
365 let start = Instant::now();
366 let result = self
367 .handle_fetch(
368 &p,
369 &correlation_id,
370 call.caller_id.clone(),
371 call.skill_name.clone(),
372 )
373 .await;
374 #[allow(clippy::cast_possible_truncation)]
375 let duration_ms = start.elapsed().as_millis() as u64;
376 self.run_with_audit(
377 "fetch",
378 "fetch",
379 &redact_url_for_log(&p.url),
380 call.caller_id.clone(),
381 call.skill_name.clone(),
382 correlation_id,
383 duration_ms,
384 result,
385 )
386 .await
387 }
388 _ => Ok(None),
389 }
390 }
391
392 fn is_tool_retryable(&self, tool_id: &str) -> bool {
393 matches!(tool_id, "web_scrape" | "fetch")
394 }
395
396 crate::tool_executor_no_inner_defaults!();
397}
398
399fn tool_error_to_audit_result(e: &ToolError) -> AuditResult {
400 match e {
401 ToolError::Blocked { command } => AuditResult::Blocked {
402 reason: command.clone(),
403 },
404 ToolError::Timeout { .. } => AuditResult::Timeout,
405 _ => AuditResult::Error {
406 message: e.to_string(),
407 },
408 }
409}
410
411impl WebScrapeExecutor {
412 #[allow(clippy::too_many_arguments)]
413 async fn run_with_audit(
414 &self,
415 audit_tool_name: &str,
416 public_tool_name: &str,
417 audit_command: &str,
418 caller_id: Option<String>,
419 skill_name: Option<Vec<String>>,
420 correlation_id: String,
421 duration_ms: u64,
422 result: Result<String, ToolError>,
423 ) -> Result<Option<ToolOutput>, ToolError> {
424 match result {
425 Ok(output) => {
426 self.log_audit(
427 audit_tool_name,
428 audit_command,
429 AuditResult::Success,
430 duration_ms,
431 None,
432 caller_id,
433 skill_name,
434 Some(correlation_id),
435 )
436 .await;
437 Ok(Some(ToolOutput {
438 tool_name: ToolName::new(public_tool_name),
439 summary: output,
440 blocks_executed: 1,
441 filter_stats: None,
442 diff: None,
443 streamed: false,
444 terminal_id: None,
445 locations: None,
446 raw_response: None,
447 claim_source: Some(ClaimSource::WebScrape),
448 ..Default::default()
449 }))
450 }
451 Err(e) => {
452 let audit_result = tool_error_to_audit_result(&e);
453 self.log_audit(
454 audit_tool_name,
455 audit_command,
456 audit_result,
457 duration_ms,
458 Some(&e),
459 caller_id,
460 skill_name,
461 Some(correlation_id),
462 )
463 .await;
464 Err(e)
465 }
466 }
467 }
468
469 #[allow(clippy::too_many_arguments)] async fn log_audit(
471 &self,
472 tool: &str,
473 command: &str,
474 result: AuditResult,
475 duration_ms: u64,
476 error: Option<&ToolError>,
477 caller_id: Option<String>,
478 skill_name: Option<Vec<String>>,
479 correlation_id: Option<String>,
480 ) {
481 if let Some(ref logger) = self.audit_logger {
482 let (error_category, error_domain, error_phase) =
483 error.map_or((None, None, None), |e| {
484 let cat = e.category();
485 (
486 Some(cat.label().to_owned()),
487 Some(cat.domain().label().to_owned()),
488 Some(cat.phase().label().to_owned()),
489 )
490 });
491 let entry = AuditEntry {
492 source_kind: None,
493 trust_level: None,
494 timestamp: chrono_now(),
495 tool: tool.into(),
496 command: command.into(),
497 result,
498 duration_ms,
499 error_category,
500 error_domain,
501 error_phase,
502 claim_source: Some(ClaimSource::WebScrape),
503 mcp_server_id: None,
504 injection_flagged: false,
505 embedding_anomalous: false,
506 cross_boundary_mcp_to_acp: false,
507 adversarial_policy_decision: None,
508 exit_code: None,
509 truncated: false,
510 caller_id,
511 skill_name,
512 policy_match: None,
513 correlation_id,
514 vigil_risk: None,
515 execution_env: None,
516 resolved_cwd: None,
517 scope_at_definition: None,
518 scope_at_dispatch: None,
519 };
520 logger.log(&entry).await;
521 }
522 }
523
524 fn send_egress_event(&self, event: EgressEvent) {
525 if let Some(ref tx) = self.egress_tx {
526 match tx.try_send(event) {
527 Ok(()) => {}
528 Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
529 self.egress_dropped.fetch_add(1, Ordering::Relaxed);
530 }
531 Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
532 tracing::debug!("egress channel closed; executor continuing without telemetry");
533 }
534 }
535 }
536 }
537
538 async fn log_egress_event(&self, event: &EgressEvent) {
539 if let Some(ref logger) = self.audit_logger {
540 logger.log_egress(event).await;
541 }
542 self.send_egress_event(event.clone());
543 }
544
545 async fn handle_fetch(
546 &self,
547 params: &FetchParams,
548 correlation_id: &str,
549 caller_id: Option<String>,
550 skill_name: Option<Vec<String>>,
551 ) -> Result<String, ToolError> {
552 let parsed = validate_url(¶ms.url);
553 let host_str = parsed
554 .as_ref()
555 .map(|u| u.host_str().unwrap_or("").to_owned())
556 .unwrap_or_default();
557
558 if let Err(ref _e) = parsed {
559 if self.egress_config.enabled && self.egress_config.log_blocked {
560 let event = Self::make_blocked_event(
561 "fetch",
562 ¶ms.url,
563 &host_str,
564 correlation_id,
565 caller_id.clone(),
566 skill_name.clone(),
567 "scheme",
568 );
569 self.log_egress_event(&event).await;
570 }
571 return Err(parsed.unwrap_err());
572 }
573 let parsed = parsed.unwrap();
574
575 if let Err(e) = check_domain_policy(
576 parsed.host_str().unwrap_or(""),
577 &self.allowed_domains,
578 &self.denied_domains,
579 ) {
580 if self.egress_config.enabled && self.egress_config.log_blocked {
581 let event = Self::make_blocked_event(
582 "fetch",
583 ¶ms.url,
584 parsed.host_str().unwrap_or(""),
585 correlation_id,
586 caller_id.clone(),
587 skill_name.clone(),
588 "blocklist",
589 );
590 self.log_egress_event(&event).await;
591 }
592 return Err(e);
593 }
594
595 let (host, addrs) = match resolve_and_validate(&parsed).await {
596 Ok(v) => v,
597 Err(e) => {
598 if self.egress_config.enabled && self.egress_config.log_blocked {
599 let event = Self::make_blocked_event(
600 "fetch",
601 ¶ms.url,
602 parsed.host_str().unwrap_or(""),
603 correlation_id,
604 caller_id.clone(),
605 skill_name.clone(),
606 "ssrf",
607 );
608 self.log_egress_event(&event).await;
609 }
610 return Err(e);
611 }
612 };
613
614 let body = self
615 .fetch_html(
616 ¶ms.url,
617 &host,
618 &addrs,
619 "fetch",
620 correlation_id,
621 caller_id,
622 skill_name,
623 )
624 .await?;
625 self.apply_ipi_filter(&body, ¶ms.url).await
626 }
627
628 async fn scrape_instruction(
629 &self,
630 instruction: &ScrapeInstruction,
631 correlation_id: &str,
632 caller_id: Option<String>,
633 skill_name: Option<Vec<String>>,
634 ) -> Result<String, ToolError> {
635 let parsed = validate_url(&instruction.url);
636 let host_str = parsed
637 .as_ref()
638 .map(|u| u.host_str().unwrap_or("").to_owned())
639 .unwrap_or_default();
640
641 if let Err(ref _e) = parsed {
642 if self.egress_config.enabled && self.egress_config.log_blocked {
643 let event = Self::make_blocked_event(
644 "web_scrape",
645 &instruction.url,
646 &host_str,
647 correlation_id,
648 caller_id.clone(),
649 skill_name.clone(),
650 "scheme",
651 );
652 self.log_egress_event(&event).await;
653 }
654 return Err(parsed.unwrap_err());
655 }
656 let parsed = parsed.unwrap();
657
658 if let Err(e) = check_domain_policy(
659 parsed.host_str().unwrap_or(""),
660 &self.allowed_domains,
661 &self.denied_domains,
662 ) {
663 if self.egress_config.enabled && self.egress_config.log_blocked {
664 let event = Self::make_blocked_event(
665 "web_scrape",
666 &instruction.url,
667 parsed.host_str().unwrap_or(""),
668 correlation_id,
669 caller_id.clone(),
670 skill_name.clone(),
671 "blocklist",
672 );
673 self.log_egress_event(&event).await;
674 }
675 return Err(e);
676 }
677
678 let (host, addrs) = match resolve_and_validate(&parsed).await {
679 Ok(v) => v,
680 Err(e) => {
681 if self.egress_config.enabled && self.egress_config.log_blocked {
682 let event = Self::make_blocked_event(
683 "web_scrape",
684 &instruction.url,
685 parsed.host_str().unwrap_or(""),
686 correlation_id,
687 caller_id.clone(),
688 skill_name.clone(),
689 "ssrf",
690 );
691 self.log_egress_event(&event).await;
692 }
693 return Err(e);
694 }
695 };
696
697 let html = self
698 .fetch_html(
699 &instruction.url,
700 &host,
701 &addrs,
702 "web_scrape",
703 correlation_id,
704 caller_id,
705 skill_name,
706 )
707 .await?;
708 let selector = instruction.select.clone();
709 let extract = ExtractMode::parse(&instruction.extract);
710 let limit = instruction.limit.unwrap_or(10);
711 let extracted = tokio::task::spawn_blocking(move || {
712 parse_and_extract(&html, &selector, &extract, limit)
713 })
714 .await
715 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))??;
716 self.apply_ipi_filter(&extracted, &instruction.url).await
718 }
719
720 fn make_blocked_event(
721 tool: &str,
722 url: &str,
723 host: &str,
724 correlation_id: &str,
725 caller_id: Option<String>,
726 skill_name: Option<Vec<String>>,
727 block_reason: &'static str,
728 ) -> EgressEvent {
729 EgressEvent {
730 timestamp: chrono_now(),
731 kind: "egress",
732 correlation_id: correlation_id.to_owned(),
733 tool: tool.into(),
734 url: redact_url_for_log(url),
735 host: host.to_owned(),
736 method: "GET".to_owned(),
737 status: None,
738 duration_ms: 0,
739 response_bytes: 0,
740 blocked: true,
741 block_reason: Some(block_reason),
742 caller_id,
743 skill_name,
744 hop: 0,
745 }
746 }
747
748 #[allow(clippy::too_many_lines, clippy::too_many_arguments)]
759 async fn fetch_html(
760 &self,
761 url: &str,
762 host: &str,
763 addrs: &[SocketAddr],
764 tool: &str,
765 correlation_id: &str,
766 caller_id: Option<String>,
767 skill_name: Option<Vec<String>>,
768 ) -> Result<String, ToolError> {
769 const MAX_REDIRECTS: usize = 3;
770
771 let mut current_url = url.to_owned();
772 let mut current_host = host.to_owned();
773 let mut current_addrs = addrs.to_vec();
774
775 for hop in 0..=MAX_REDIRECTS {
776 let hop_start = Instant::now();
777 let client = self.build_client(¤t_host, ¤t_addrs);
779 let resp = client
780 .get(¤t_url)
781 .send()
782 .await
783 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())));
784
785 let resp = match resp {
786 Ok(r) => r,
787 Err(e) => {
788 if self.egress_config.enabled {
789 #[allow(clippy::cast_possible_truncation)]
790 let duration_ms = hop_start.elapsed().as_millis() as u64;
791 let event = EgressEvent {
792 timestamp: chrono_now(),
793 kind: "egress",
794 correlation_id: correlation_id.to_owned(),
795 tool: tool.into(),
796 url: redact_url_for_log(¤t_url),
797 host: current_host.clone(),
798 method: "GET".to_owned(),
799 status: None,
800 duration_ms,
801 response_bytes: 0,
802 blocked: false,
803 block_reason: None,
804 caller_id: caller_id.clone(),
805 skill_name: skill_name.clone(),
806 #[allow(clippy::cast_possible_truncation)]
807 hop: hop as u8,
808 };
809 self.log_egress_event(&event).await;
810 }
811 return Err(e);
812 }
813 };
814
815 let status = resp.status();
816
817 if status.is_redirection() {
818 if hop == MAX_REDIRECTS {
819 return Err(ToolError::Execution(std::io::Error::other(
820 "too many redirects",
821 )));
822 }
823
824 let location = resp
825 .headers()
826 .get(reqwest::header::LOCATION)
827 .and_then(|v| v.to_str().ok())
828 .ok_or_else(|| {
829 ToolError::Execution(std::io::Error::other("redirect with no Location"))
830 })?;
831
832 let base = Url::parse(¤t_url)
834 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
835 let next_url = base
836 .join(location)
837 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
838
839 let validated = validate_url(next_url.as_str());
840 if let Err(ref _e) = validated {
841 if self.egress_config.enabled && self.egress_config.log_blocked {
842 #[allow(clippy::cast_possible_truncation)]
843 let duration_ms = hop_start.elapsed().as_millis() as u64;
844 let next_host = next_url.host_str().unwrap_or("").to_owned();
845 let event = EgressEvent {
846 timestamp: chrono_now(),
847 kind: "egress",
848 correlation_id: correlation_id.to_owned(),
849 tool: tool.into(),
850 url: redact_url_for_log(next_url.as_str()),
851 host: next_host,
852 method: "GET".to_owned(),
853 status: None,
854 duration_ms,
855 response_bytes: 0,
856 blocked: true,
857 block_reason: Some("ssrf"),
858 caller_id: caller_id.clone(),
859 skill_name: skill_name.clone(),
860 #[allow(clippy::cast_possible_truncation)]
861 hop: (hop + 1) as u8,
862 };
863 self.log_egress_event(&event).await;
864 }
865 return Err(validated.unwrap_err());
866 }
867 let validated = validated.unwrap();
868 let resolve_result = resolve_and_validate(&validated).await;
869 if let Err(ref _e) = resolve_result {
870 if self.egress_config.enabled && self.egress_config.log_blocked {
871 #[allow(clippy::cast_possible_truncation)]
872 let duration_ms = hop_start.elapsed().as_millis() as u64;
873 let next_host = next_url.host_str().unwrap_or("").to_owned();
874 let event = EgressEvent {
875 timestamp: chrono_now(),
876 kind: "egress",
877 correlation_id: correlation_id.to_owned(),
878 tool: tool.into(),
879 url: redact_url_for_log(next_url.as_str()),
880 host: next_host,
881 method: "GET".to_owned(),
882 status: None,
883 duration_ms,
884 response_bytes: 0,
885 blocked: true,
886 block_reason: Some("ssrf"),
887 caller_id: caller_id.clone(),
888 skill_name: skill_name.clone(),
889 #[allow(clippy::cast_possible_truncation)]
890 hop: (hop + 1) as u8,
891 };
892 self.log_egress_event(&event).await;
893 }
894 return Err(resolve_result.unwrap_err());
895 }
896 let (next_host, next_addrs) = resolve_result.unwrap();
897
898 current_url = next_url.to_string();
899 current_host = next_host;
900 current_addrs = next_addrs;
901 continue;
902 }
903
904 if !status.is_success() {
905 if self.egress_config.enabled {
906 #[allow(clippy::cast_possible_truncation)]
907 let duration_ms = hop_start.elapsed().as_millis() as u64;
908 let event = EgressEvent {
909 timestamp: chrono_now(),
910 kind: "egress",
911 correlation_id: correlation_id.to_owned(),
912 tool: tool.into(),
913 url: redact_url_for_log(¤t_url),
914 host: current_host.clone(),
915 method: "GET".to_owned(),
916 status: Some(status.as_u16()),
917 duration_ms,
918 response_bytes: 0,
919 blocked: false,
920 block_reason: None,
921 caller_id: caller_id.clone(),
922 skill_name: skill_name.clone(),
923 #[allow(clippy::cast_possible_truncation)]
924 hop: hop as u8,
925 };
926 self.log_egress_event(&event).await;
927 }
928 return Err(ToolError::Http {
929 status: status.as_u16(),
930 message: status.canonical_reason().unwrap_or("unknown").to_owned(),
931 });
932 }
933
934 let bytes = resp
935 .bytes()
936 .await
937 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
938
939 if bytes.len() > self.max_body_bytes {
940 if self.egress_config.enabled {
941 #[allow(clippy::cast_possible_truncation)]
942 let duration_ms = hop_start.elapsed().as_millis() as u64;
943 let event = EgressEvent {
944 timestamp: chrono_now(),
945 kind: "egress",
946 correlation_id: correlation_id.to_owned(),
947 tool: tool.into(),
948 url: redact_url_for_log(¤t_url),
949 host: current_host.clone(),
950 method: "GET".to_owned(),
951 status: Some(status.as_u16()),
952 duration_ms,
953 response_bytes: bytes.len(),
954 blocked: false,
955 block_reason: None,
956 caller_id: caller_id.clone(),
957 skill_name: skill_name.clone(),
958 #[allow(clippy::cast_possible_truncation)]
959 hop: hop as u8,
960 };
961 self.log_egress_event(&event).await;
962 }
963 return Err(ToolError::Execution(std::io::Error::other(format!(
964 "response too large: {} bytes (max: {})",
965 bytes.len(),
966 self.max_body_bytes,
967 ))));
968 }
969
970 if self.egress_config.enabled {
972 #[allow(clippy::cast_possible_truncation)]
973 let duration_ms = hop_start.elapsed().as_millis() as u64;
974 let response_bytes = if self.egress_config.log_response_bytes {
975 bytes.len()
976 } else {
977 0
978 };
979 let event = EgressEvent {
980 timestamp: chrono_now(),
981 kind: "egress",
982 correlation_id: correlation_id.to_owned(),
983 tool: tool.into(),
984 url: redact_url_for_log(¤t_url),
985 host: current_host.clone(),
986 method: "GET".to_owned(),
987 status: Some(status.as_u16()),
988 duration_ms,
989 response_bytes,
990 blocked: false,
991 block_reason: None,
992 caller_id: caller_id.clone(),
993 skill_name: skill_name.clone(),
994 #[allow(clippy::cast_possible_truncation)]
995 hop: hop as u8,
996 };
997 self.log_egress_event(&event).await;
998 }
999
1000 return String::from_utf8(bytes.to_vec())
1001 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())));
1002 }
1003
1004 Err(ToolError::Execution(std::io::Error::other(
1005 "too many redirects",
1006 )))
1007 }
1008
1009 #[tracing::instrument(name = "tools.scrape.apply_ipi_filter", skip(self, body), fields(body_len = body.len()))]
1019 async fn apply_ipi_filter(&self, body: &str, url: &str) -> Result<String, ToolError> {
1020 let verdict = self
1021 .ipi_filter
1022 .filter_async(body.to_owned())
1023 .await
1024 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
1025 if !verdict.patterns_found.is_empty() {
1026 tracing::warn!(
1027 url = url,
1028 score = verdict.score,
1029 patterns = ?verdict.patterns_found,
1030 "IPI patterns detected in fetched content"
1031 );
1032 }
1033 if verdict.sanitized == body {
1035 Ok(verdict.sanitized)
1036 } else {
1037 Ok(format!(
1038 "[IPI WARNING: score={:.2}, patterns={}] {}",
1039 verdict.score,
1040 verdict.patterns_found.join(", "),
1041 verdict.sanitized,
1042 ))
1043 }
1044 }
1045}
1046
1047fn extract_scrape_blocks(text: &str) -> Vec<&str> {
1048 crate::executor::extract_fenced_blocks(text, "scrape")
1049}
1050
1051#[tracing::instrument(name = "tools.scrape.dns.resolve", skip(url), fields(host = url.host_str().unwrap_or("")))]
1061async fn resolve_and_validate(url: &Url) -> Result<(String, Vec<SocketAddr>), ToolError> {
1062 let Some(host) = url.host_str() else {
1063 return Ok((String::new(), vec![]));
1064 };
1065 let port = url.port_or_known_default().unwrap_or(443);
1066 let addrs = zeph_common::net::resolve_and_validate(host, port)
1067 .await
1068 .map_err(|e| match e {
1069 zeph_common::net::ResolveError::Timeout(timeout) => ToolError::Timeout {
1070 timeout_secs: timeout.as_secs(),
1071 },
1072 zeph_common::net::ResolveError::Lookup(io_err) => ToolError::Blocked {
1073 command: format!("DNS resolution failed: {io_err}"),
1074 },
1075 zeph_common::net::ResolveError::PrivateAddress { host, addr } => ToolError::Blocked {
1076 command: format!("SSRF protection: private IP {addr} for host {host}"),
1077 },
1078 other => ToolError::Blocked {
1079 command: format!("DNS resolution failed: {other}"),
1080 },
1081 })?;
1082 Ok((host.to_owned(), addrs))
1083}
1084
1085fn parse_and_extract(
1086 html: &str,
1087 selector: &str,
1088 extract: &ExtractMode,
1089 limit: usize,
1090) -> Result<String, ToolError> {
1091 let soup = scrape_core::Soup::parse(html);
1092
1093 let tags = soup.find_all(selector).map_err(|e| {
1094 ToolError::Execution(std::io::Error::new(
1095 std::io::ErrorKind::InvalidData,
1096 format!("invalid selector: {e}"),
1097 ))
1098 })?;
1099
1100 let mut results = Vec::new();
1101
1102 for tag in tags.into_iter().take(limit) {
1103 let value = match extract {
1104 ExtractMode::Text => tag.text(),
1105 ExtractMode::Html => tag.inner_html(),
1106 ExtractMode::Attr(name) => tag.get(name).unwrap_or_default().to_owned(),
1107 };
1108 if !value.trim().is_empty() {
1109 results.push(value.trim().to_owned());
1110 }
1111 }
1112
1113 if results.is_empty() {
1114 Ok(format!("No results for selector: {selector}"))
1115 } else {
1116 Ok(results.join("\n"))
1117 }
1118}
1119
1120#[cfg(test)]
1121mod tests {
1122 use super::*;
1123 use crate::domain_match::domain_matches;
1124 use crate::net::is_private_host;
1125 use std::assert_matches;
1126
1127 #[test]
1130 fn extract_single_block() {
1131 let text =
1132 "Here:\n```scrape\n{\"url\":\"https://example.com\",\"select\":\"h1\"}\n```\nDone.";
1133 let blocks = extract_scrape_blocks(text);
1134 assert_eq!(blocks.len(), 1);
1135 assert!(blocks[0].contains("example.com"));
1136 }
1137
1138 #[test]
1139 fn extract_multiple_blocks() {
1140 let text = "```scrape\n{\"url\":\"https://a.com\",\"select\":\"h1\"}\n```\ntext\n```scrape\n{\"url\":\"https://b.com\",\"select\":\"p\"}\n```";
1141 let blocks = extract_scrape_blocks(text);
1142 assert_eq!(blocks.len(), 2);
1143 }
1144
1145 #[test]
1146 fn no_blocks_returns_empty() {
1147 let blocks = extract_scrape_blocks("plain text, no code blocks");
1148 assert!(blocks.is_empty());
1149 }
1150
1151 #[test]
1152 fn unclosed_block_ignored() {
1153 let blocks = extract_scrape_blocks("```scrape\n{\"url\":\"https://x.com\"}");
1154 assert!(blocks.is_empty());
1155 }
1156
1157 #[test]
1158 fn non_scrape_block_ignored() {
1159 let text =
1160 "```bash\necho hi\n```\n```scrape\n{\"url\":\"https://x.com\",\"select\":\"h1\"}\n```";
1161 let blocks = extract_scrape_blocks(text);
1162 assert_eq!(blocks.len(), 1);
1163 assert!(blocks[0].contains("x.com"));
1164 }
1165
1166 #[test]
1167 fn multiline_json_block() {
1168 let text =
1169 "```scrape\n{\n \"url\": \"https://example.com\",\n \"select\": \"h1\"\n}\n```";
1170 let blocks = extract_scrape_blocks(text);
1171 assert_eq!(blocks.len(), 1);
1172 let instr: ScrapeInstruction = serde_json::from_str(blocks[0]).unwrap();
1173 assert_eq!(instr.url, "https://example.com");
1174 }
1175
1176 #[test]
1179 fn parse_valid_instruction() {
1180 let json = r#"{"url":"https://example.com","select":"h1","extract":"text","limit":5}"#;
1181 let instr: ScrapeInstruction = serde_json::from_str(json).unwrap();
1182 assert_eq!(instr.url, "https://example.com");
1183 assert_eq!(instr.select, "h1");
1184 assert_eq!(instr.extract, "text");
1185 assert_eq!(instr.limit, Some(5));
1186 }
1187
1188 #[test]
1189 fn parse_minimal_instruction() {
1190 let json = r#"{"url":"https://example.com","select":"p"}"#;
1191 let instr: ScrapeInstruction = serde_json::from_str(json).unwrap();
1192 assert_eq!(instr.extract, "text");
1193 assert!(instr.limit.is_none());
1194 }
1195
1196 #[test]
1197 fn parse_attr_extract() {
1198 let json = r#"{"url":"https://example.com","select":"a","extract":"attr:href"}"#;
1199 let instr: ScrapeInstruction = serde_json::from_str(json).unwrap();
1200 assert_eq!(instr.extract, "attr:href");
1201 }
1202
1203 #[test]
1204 fn parse_invalid_json_errors() {
1205 let result = serde_json::from_str::<ScrapeInstruction>("not json");
1206 assert!(result.is_err());
1207 }
1208
1209 #[test]
1212 fn extract_mode_text() {
1213 assert_matches!(ExtractMode::parse("text"), ExtractMode::Text);
1214 }
1215
1216 #[test]
1217 fn extract_mode_html() {
1218 assert_matches!(ExtractMode::parse("html"), ExtractMode::Html);
1219 }
1220
1221 #[test]
1222 fn extract_mode_attr() {
1223 let mode = ExtractMode::parse("attr:href");
1224 assert_matches!(mode, ExtractMode::Attr(ref s) if s == "href");
1225 }
1226
1227 #[test]
1228 fn extract_mode_unknown_defaults_to_text() {
1229 assert_matches!(ExtractMode::parse("unknown"), ExtractMode::Text);
1230 }
1231
1232 #[test]
1238 fn extract_text_from_html() {
1239 let html = "<html><body><h1>Hello World</h1><p>Content</p></body></html>";
1240 let result = parse_and_extract(html, "h1", &ExtractMode::Text, 10).unwrap();
1241 assert_eq!(result, "Hello World");
1242 }
1243
1244 #[test]
1245 fn extract_multiple_elements() {
1246 let html = "<ul><li>A</li><li>B</li><li>C</li></ul>";
1247 let result = parse_and_extract(html, "li", &ExtractMode::Text, 10).unwrap();
1248 assert_eq!(result, "A\nB\nC");
1249 }
1250
1251 #[test]
1252 fn extract_with_limit() {
1253 let html = "<ul><li>A</li><li>B</li><li>C</li></ul>";
1254 let result = parse_and_extract(html, "li", &ExtractMode::Text, 2).unwrap();
1255 assert_eq!(result, "A\nB");
1256 }
1257
1258 #[test]
1259 fn extract_attr_href() {
1260 let html = r#"<a href="https://example.com">Link</a>"#;
1261 let result =
1262 parse_and_extract(html, "a", &ExtractMode::Attr("href".to_owned()), 10).unwrap();
1263 assert_eq!(result, "https://example.com");
1264 }
1265
1266 #[test]
1267 fn extract_inner_html() {
1268 let html = "<div><span>inner</span></div>";
1269 let result = parse_and_extract(html, "div", &ExtractMode::Html, 10).unwrap();
1270 assert!(result.contains("<span>inner</span>"));
1271 }
1272
1273 #[test]
1274 fn no_matches_returns_message() {
1275 let html = "<html><body><p>text</p></body></html>";
1276 let result = parse_and_extract(html, "h1", &ExtractMode::Text, 10).unwrap();
1277 assert!(result.starts_with("No results for selector:"));
1278 }
1279
1280 #[test]
1281 fn empty_text_skipped() {
1282 let html = "<ul><li> </li><li>A</li></ul>";
1283 let result = parse_and_extract(html, "li", &ExtractMode::Text, 10).unwrap();
1284 assert_eq!(result, "A");
1285 }
1286
1287 #[test]
1288 fn invalid_selector_errors() {
1289 let html = "<html><body></body></html>";
1290 let result = parse_and_extract(html, "[[[invalid", &ExtractMode::Text, 10);
1291 assert!(result.is_err());
1292 }
1293
1294 #[test]
1295 fn empty_html_returns_no_results() {
1296 let result = parse_and_extract("", "h1", &ExtractMode::Text, 10).unwrap();
1297 assert!(result.starts_with("No results for selector:"));
1298 }
1299
1300 #[test]
1301 fn nested_selector() {
1302 let html = "<div><span>inner</span></div><span>outer</span>";
1303 let result = parse_and_extract(html, "div > span", &ExtractMode::Text, 10).unwrap();
1304 assert_eq!(result, "inner");
1305 }
1306
1307 #[test]
1308 fn attr_missing_returns_empty() {
1309 let html = r"<a>No href</a>";
1310 let result =
1311 parse_and_extract(html, "a", &ExtractMode::Attr("href".to_owned()), 10).unwrap();
1312 assert!(result.starts_with("No results for selector:"));
1313 }
1314
1315 #[test]
1316 fn extract_html_mode() {
1317 let html = "<div><b>bold</b> text</div>";
1318 let result = parse_and_extract(html, "div", &ExtractMode::Html, 10).unwrap();
1319 assert!(result.contains("<b>bold</b>"));
1320 }
1321
1322 #[test]
1323 fn limit_zero_returns_no_results() {
1324 let html = "<ul><li>A</li><li>B</li></ul>";
1325 let result = parse_and_extract(html, "li", &ExtractMode::Text, 0).unwrap();
1326 assert!(result.starts_with("No results for selector:"));
1327 }
1328
1329 #[tokio::test]
1332 async fn executor_no_blocks_returns_none() {
1333 let config = ScrapeConfig::default();
1334 let executor = WebScrapeExecutor::new(&config);
1335 let result = executor.execute("plain text").await;
1336 assert!(result.unwrap().is_none());
1337 }
1338
1339 #[tokio::test]
1340 async fn executor_invalid_json_errors() {
1341 let config = ScrapeConfig::default();
1342 let executor = WebScrapeExecutor::new(&config);
1343 let response = "```scrape\nnot json\n```";
1344 let result = executor.execute(response).await;
1345 assert_matches!(result, Err(ToolError::Execution(_)));
1346 }
1347
1348 #[tokio::test]
1349 async fn executor_blocked_url_errors() {
1350 let config = ScrapeConfig::default();
1351 let executor = WebScrapeExecutor::new(&config);
1352 let response = "```scrape\n{\"url\":\"http://example.com\",\"select\":\"h1\"}\n```";
1353 let result = executor.execute(response).await;
1354 assert_matches!(result, Err(ToolError::Blocked { .. }));
1355 }
1356
1357 #[tokio::test]
1358 async fn executor_private_ip_blocked() {
1359 let config = ScrapeConfig::default();
1360 let executor = WebScrapeExecutor::new(&config);
1361 let response = "```scrape\n{\"url\":\"https://192.168.1.1/api\",\"select\":\"h1\"}\n```";
1362 let result = executor.execute(response).await;
1363 assert_matches!(result, Err(ToolError::Blocked { .. }));
1364 }
1365
1366 #[tokio::test]
1367 async fn executor_unreachable_host_returns_error() {
1368 let config = ScrapeConfig {
1369 timeout: 1,
1370 max_body_bytes: 1_048_576,
1371 ..Default::default()
1372 };
1373 let executor = WebScrapeExecutor::new(&config);
1374 let response = "```scrape\n{\"url\":\"https://192.0.2.1:1/page\",\"select\":\"h1\"}\n```";
1375 let result = executor.execute(response).await;
1376 assert_matches!(result, Err(ToolError::Execution(_)));
1377 }
1378
1379 #[tokio::test]
1380 async fn executor_localhost_url_blocked() {
1381 let config = ScrapeConfig::default();
1382 let executor = WebScrapeExecutor::new(&config);
1383 let response = "```scrape\n{\"url\":\"https://localhost:9999/api\",\"select\":\"h1\"}\n```";
1384 let result = executor.execute(response).await;
1385 assert_matches!(result, Err(ToolError::Blocked { .. }));
1386 }
1387
1388 #[tokio::test]
1389 async fn executor_empty_text_returns_none() {
1390 let config = ScrapeConfig::default();
1391 let executor = WebScrapeExecutor::new(&config);
1392 let result = executor.execute("").await;
1393 assert!(result.unwrap().is_none());
1394 }
1395
1396 #[tokio::test]
1397 async fn executor_multiple_blocks_first_blocked() {
1398 let config = ScrapeConfig::default();
1399 let executor = WebScrapeExecutor::new(&config);
1400 let response = "```scrape\n{\"url\":\"http://evil.com\",\"select\":\"h1\"}\n```\n\
1401 ```scrape\n{\"url\":\"https://ok.com\",\"select\":\"h1\"}\n```";
1402 let result = executor.execute(response).await;
1403 assert!(result.is_err());
1404 }
1405
1406 #[test]
1407 fn validate_url_empty_string() {
1408 let err = validate_url("").unwrap_err();
1409 assert_matches!(err, ToolError::Blocked { .. });
1410 }
1411
1412 #[test]
1413 fn validate_url_javascript_scheme_blocked() {
1414 let err = validate_url("javascript:alert(1)").unwrap_err();
1415 assert_matches!(err, ToolError::Blocked { .. });
1416 }
1417
1418 #[test]
1419 fn validate_url_data_scheme_blocked() {
1420 let err = validate_url("data:text/html,<h1>hi</h1>").unwrap_err();
1421 assert_matches!(err, ToolError::Blocked { .. });
1422 }
1423
1424 #[test]
1425 fn is_private_host_public_domain_is_false() {
1426 let host: url::Host<&str> = url::Host::Domain("example.com");
1427 assert!(!is_private_host(&host));
1428 }
1429
1430 #[test]
1431 fn is_private_host_localhost_is_true() {
1432 let host: url::Host<&str> = url::Host::Domain("localhost");
1433 assert!(is_private_host(&host));
1434 }
1435
1436 #[test]
1437 fn is_private_host_ipv6_unspecified_is_true() {
1438 let host = url::Host::Ipv6(std::net::Ipv6Addr::UNSPECIFIED);
1439 assert!(is_private_host(&host));
1440 }
1441
1442 #[test]
1443 fn is_private_host_public_ipv6_is_false() {
1444 let host = url::Host::Ipv6("2001:db8::1".parse().unwrap());
1445 assert!(!is_private_host(&host));
1446 }
1447
1448 async fn mock_server_executor() -> (WebScrapeExecutor, wiremock::MockServer) {
1459 let server = wiremock::MockServer::start().await;
1460 let executor = WebScrapeExecutor {
1461 timeout: Duration::from_secs(5),
1462 max_body_bytes: 1_048_576,
1463 allowed_domains: vec![],
1464 denied_domains: vec![],
1465 audit_logger: None,
1466 egress_config: EgressConfig::default(),
1467 egress_tx: None,
1468 egress_dropped: Arc::new(AtomicU64::new(0)),
1469 ipi_filter: IpiFilter::new(0.6),
1470 };
1471 (executor, server)
1472 }
1473
1474 fn server_host_and_addr(server: &wiremock::MockServer) -> (String, Vec<std::net::SocketAddr>) {
1476 let uri = server.uri();
1477 let url = Url::parse(&uri).unwrap();
1478 let host = url.host_str().unwrap_or("127.0.0.1").to_owned();
1479 let port = url.port().unwrap_or(80);
1480 let addr: std::net::SocketAddr = format!("{host}:{port}").parse().unwrap();
1481 (host, vec![addr])
1482 }
1483
1484 async fn follow_redirects_raw(
1488 executor: &WebScrapeExecutor,
1489 start_url: &str,
1490 host: &str,
1491 addrs: &[std::net::SocketAddr],
1492 ) -> Result<String, ToolError> {
1493 const MAX_REDIRECTS: usize = 3;
1494 let mut current_url = start_url.to_owned();
1495 let mut current_host = host.to_owned();
1496 let mut current_addrs = addrs.to_vec();
1497
1498 for hop in 0..=MAX_REDIRECTS {
1499 let client = executor.build_client(¤t_host, ¤t_addrs);
1500 let resp = client
1501 .get(¤t_url)
1502 .send()
1503 .await
1504 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
1505
1506 let status = resp.status();
1507
1508 if status.is_redirection() {
1509 if hop == MAX_REDIRECTS {
1510 return Err(ToolError::Execution(std::io::Error::other(
1511 "too many redirects",
1512 )));
1513 }
1514
1515 let location = resp
1516 .headers()
1517 .get(reqwest::header::LOCATION)
1518 .and_then(|v| v.to_str().ok())
1519 .ok_or_else(|| {
1520 ToolError::Execution(std::io::Error::other("redirect with no Location"))
1521 })?;
1522
1523 let base = Url::parse(¤t_url)
1524 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
1525 let next_url = base
1526 .join(location)
1527 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
1528
1529 current_url = next_url.to_string();
1531 let _ = &mut current_host;
1533 let _ = &mut current_addrs;
1534 continue;
1535 }
1536
1537 if !status.is_success() {
1538 return Err(ToolError::Execution(std::io::Error::other(format!(
1539 "HTTP {status}",
1540 ))));
1541 }
1542
1543 let bytes = resp
1544 .bytes()
1545 .await
1546 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())))?;
1547
1548 if bytes.len() > executor.max_body_bytes {
1549 return Err(ToolError::Execution(std::io::Error::other(format!(
1550 "response too large: {} bytes (max: {})",
1551 bytes.len(),
1552 executor.max_body_bytes,
1553 ))));
1554 }
1555
1556 return String::from_utf8(bytes.to_vec())
1557 .map_err(|e| ToolError::Execution(std::io::Error::other(e.to_string())));
1558 }
1559
1560 Err(ToolError::Execution(std::io::Error::other(
1561 "too many redirects",
1562 )))
1563 }
1564
1565 #[tokio::test]
1566 async fn fetch_html_success_returns_body() {
1567 use wiremock::matchers::{method, path};
1568 use wiremock::{Mock, ResponseTemplate};
1569
1570 let (executor, server) = mock_server_executor().await;
1571 Mock::given(method("GET"))
1572 .and(path("/page"))
1573 .respond_with(ResponseTemplate::new(200).set_body_string("<h1>OK</h1>"))
1574 .mount(&server)
1575 .await;
1576
1577 let (host, addrs) = server_host_and_addr(&server);
1578 let url = format!("{}/page", server.uri());
1579 let result = executor
1580 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
1581 .await;
1582 assert!(result.is_ok(), "expected Ok, got: {result:?}");
1583 assert_eq!(result.unwrap(), "<h1>OK</h1>");
1584 }
1585
1586 #[tokio::test]
1587 async fn fetch_html_non_2xx_returns_error() {
1588 use wiremock::matchers::{method, path};
1589 use wiremock::{Mock, ResponseTemplate};
1590
1591 let (executor, server) = mock_server_executor().await;
1592 Mock::given(method("GET"))
1593 .and(path("/forbidden"))
1594 .respond_with(ResponseTemplate::new(403))
1595 .mount(&server)
1596 .await;
1597
1598 let (host, addrs) = server_host_and_addr(&server);
1599 let url = format!("{}/forbidden", server.uri());
1600 let result = executor
1601 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
1602 .await;
1603 assert!(result.is_err());
1604 let msg = result.unwrap_err().to_string();
1605 assert!(msg.contains("403"), "expected 403 in error: {msg}");
1606 }
1607
1608 #[tokio::test]
1609 async fn fetch_html_404_returns_error() {
1610 use wiremock::matchers::{method, path};
1611 use wiremock::{Mock, ResponseTemplate};
1612
1613 let (executor, server) = mock_server_executor().await;
1614 Mock::given(method("GET"))
1615 .and(path("/missing"))
1616 .respond_with(ResponseTemplate::new(404))
1617 .mount(&server)
1618 .await;
1619
1620 let (host, addrs) = server_host_and_addr(&server);
1621 let url = format!("{}/missing", server.uri());
1622 let result = executor
1623 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
1624 .await;
1625 assert!(result.is_err());
1626 let msg = result.unwrap_err().to_string();
1627 assert!(msg.contains("404"), "expected 404 in error: {msg}");
1628 }
1629
1630 #[tokio::test]
1631 async fn fetch_html_redirect_no_location_returns_error() {
1632 use wiremock::matchers::{method, path};
1633 use wiremock::{Mock, ResponseTemplate};
1634
1635 let (executor, server) = mock_server_executor().await;
1636 Mock::given(method("GET"))
1638 .and(path("/redirect-no-loc"))
1639 .respond_with(ResponseTemplate::new(302))
1640 .mount(&server)
1641 .await;
1642
1643 let (host, addrs) = server_host_and_addr(&server);
1644 let url = format!("{}/redirect-no-loc", server.uri());
1645 let result = executor
1646 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
1647 .await;
1648 assert!(result.is_err());
1649 let msg = result.unwrap_err().to_string();
1650 assert!(
1651 msg.contains("Location") || msg.contains("location"),
1652 "expected Location-related error: {msg}"
1653 );
1654 }
1655
1656 #[tokio::test]
1657 async fn fetch_html_single_redirect_followed() {
1658 use wiremock::matchers::{method, path};
1659 use wiremock::{Mock, ResponseTemplate};
1660
1661 let (executor, server) = mock_server_executor().await;
1662 let final_url = format!("{}/final", server.uri());
1663
1664 Mock::given(method("GET"))
1665 .and(path("/start"))
1666 .respond_with(ResponseTemplate::new(302).insert_header("location", final_url.as_str()))
1667 .mount(&server)
1668 .await;
1669
1670 Mock::given(method("GET"))
1671 .and(path("/final"))
1672 .respond_with(ResponseTemplate::new(200).set_body_string("<p>final</p>"))
1673 .mount(&server)
1674 .await;
1675
1676 let (host, addrs) = server_host_and_addr(&server);
1677 let url = format!("{}/start", server.uri());
1678 let result = follow_redirects_raw(&executor, &url, &host, &addrs).await;
1679 assert!(result.is_ok(), "single redirect should succeed: {result:?}");
1680 assert_eq!(result.unwrap(), "<p>final</p>");
1681 }
1682
1683 #[tokio::test]
1684 async fn fetch_html_three_redirects_allowed() {
1685 use wiremock::matchers::{method, path};
1686 use wiremock::{Mock, ResponseTemplate};
1687
1688 let (executor, server) = mock_server_executor().await;
1689 let hop2 = format!("{}/hop2", server.uri());
1690 let hop3 = format!("{}/hop3", server.uri());
1691 let final_dest = format!("{}/done", server.uri());
1692
1693 Mock::given(method("GET"))
1694 .and(path("/hop1"))
1695 .respond_with(ResponseTemplate::new(301).insert_header("location", hop2.as_str()))
1696 .mount(&server)
1697 .await;
1698 Mock::given(method("GET"))
1699 .and(path("/hop2"))
1700 .respond_with(ResponseTemplate::new(301).insert_header("location", hop3.as_str()))
1701 .mount(&server)
1702 .await;
1703 Mock::given(method("GET"))
1704 .and(path("/hop3"))
1705 .respond_with(ResponseTemplate::new(301).insert_header("location", final_dest.as_str()))
1706 .mount(&server)
1707 .await;
1708 Mock::given(method("GET"))
1709 .and(path("/done"))
1710 .respond_with(ResponseTemplate::new(200).set_body_string("<p>done</p>"))
1711 .mount(&server)
1712 .await;
1713
1714 let (host, addrs) = server_host_and_addr(&server);
1715 let url = format!("{}/hop1", server.uri());
1716 let result = follow_redirects_raw(&executor, &url, &host, &addrs).await;
1717 assert!(result.is_ok(), "3 redirects should succeed: {result:?}");
1718 assert_eq!(result.unwrap(), "<p>done</p>");
1719 }
1720
1721 #[tokio::test]
1722 async fn fetch_html_four_redirects_rejected() {
1723 use wiremock::matchers::{method, path};
1724 use wiremock::{Mock, ResponseTemplate};
1725
1726 let (executor, server) = mock_server_executor().await;
1727 let hop2 = format!("{}/r2", server.uri());
1728 let hop3 = format!("{}/r3", server.uri());
1729 let hop4 = format!("{}/r4", server.uri());
1730 let hop5 = format!("{}/r5", server.uri());
1731
1732 for (from, to) in [
1733 ("/r1", &hop2),
1734 ("/r2", &hop3),
1735 ("/r3", &hop4),
1736 ("/r4", &hop5),
1737 ] {
1738 Mock::given(method("GET"))
1739 .and(path(from))
1740 .respond_with(ResponseTemplate::new(301).insert_header("location", to.as_str()))
1741 .mount(&server)
1742 .await;
1743 }
1744
1745 let (host, addrs) = server_host_and_addr(&server);
1746 let url = format!("{}/r1", server.uri());
1747 let result = follow_redirects_raw(&executor, &url, &host, &addrs).await;
1748 assert!(result.is_err(), "4 redirects should be rejected");
1749 let msg = result.unwrap_err().to_string();
1750 assert!(
1751 msg.contains("redirect"),
1752 "expected redirect-related error: {msg}"
1753 );
1754 }
1755
1756 #[tokio::test]
1757 async fn fetch_html_body_too_large_returns_error() {
1758 use wiremock::matchers::{method, path};
1759 use wiremock::{Mock, ResponseTemplate};
1760
1761 let small_limit_executor = WebScrapeExecutor {
1762 timeout: Duration::from_secs(5),
1763 max_body_bytes: 10,
1764 allowed_domains: vec![],
1765 denied_domains: vec![],
1766 audit_logger: None,
1767 egress_config: EgressConfig::default(),
1768 egress_tx: None,
1769 egress_dropped: Arc::new(AtomicU64::new(0)),
1770 ipi_filter: IpiFilter::new(0.6),
1771 };
1772 let server = wiremock::MockServer::start().await;
1773 Mock::given(method("GET"))
1774 .and(path("/big"))
1775 .respond_with(
1776 ResponseTemplate::new(200)
1777 .set_body_string("this body is definitely longer than ten bytes"),
1778 )
1779 .mount(&server)
1780 .await;
1781
1782 let (host, addrs) = server_host_and_addr(&server);
1783 let url = format!("{}/big", server.uri());
1784 let result = small_limit_executor
1785 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
1786 .await;
1787 assert!(result.is_err());
1788 let msg = result.unwrap_err().to_string();
1789 assert!(msg.contains("too large"), "expected too-large error: {msg}");
1790 }
1791
1792 #[test]
1793 fn extract_scrape_blocks_empty_block_content() {
1794 let text = "```scrape\n\n```";
1795 let blocks = extract_scrape_blocks(text);
1796 assert_eq!(blocks.len(), 1);
1797 assert!(blocks[0].is_empty());
1798 }
1799
1800 #[test]
1801 fn extract_scrape_blocks_whitespace_only() {
1802 let text = "```scrape\n \n```";
1803 let blocks = extract_scrape_blocks(text);
1804 assert_eq!(blocks.len(), 1);
1805 }
1806
1807 #[test]
1808 fn parse_and_extract_multiple_selectors() {
1809 let html = "<div><h1>Title</h1><p>Para</p></div>";
1810 let result = parse_and_extract(html, "h1, p", &ExtractMode::Text, 10).unwrap();
1811 assert!(result.contains("Title"));
1812 assert!(result.contains("Para"));
1813 }
1814
1815 #[test]
1816 fn webscrape_executor_new_with_custom_config() {
1817 let config = ScrapeConfig {
1818 timeout: 60,
1819 max_body_bytes: 512,
1820 ..Default::default()
1821 };
1822 let executor = WebScrapeExecutor::new(&config);
1823 assert_eq!(executor.max_body_bytes, 512);
1824 }
1825
1826 #[test]
1827 fn webscrape_executor_debug() {
1828 let config = ScrapeConfig::default();
1829 let executor = WebScrapeExecutor::new(&config);
1830 let dbg = format!("{executor:?}");
1831 assert!(dbg.contains("WebScrapeExecutor"));
1832 }
1833
1834 #[test]
1835 fn extract_mode_attr_empty_name() {
1836 let mode = ExtractMode::parse("attr:");
1837 assert_matches!(mode, ExtractMode::Attr(ref s) if s.is_empty());
1838 }
1839
1840 #[test]
1841 fn default_extract_returns_text() {
1842 assert_eq!(default_extract(), "text");
1843 }
1844
1845 #[test]
1846 fn scrape_instruction_debug() {
1847 let json = r#"{"url":"https://example.com","select":"h1"}"#;
1848 let instr: ScrapeInstruction = serde_json::from_str(json).unwrap();
1849 let dbg = format!("{instr:?}");
1850 assert!(dbg.contains("ScrapeInstruction"));
1851 }
1852
1853 #[test]
1854 fn extract_mode_debug() {
1855 let mode = ExtractMode::Text;
1856 let dbg = format!("{mode:?}");
1857 assert!(dbg.contains("Text"));
1858 }
1859
1860 #[test]
1865 fn max_redirects_constant_is_three() {
1866 const MAX_REDIRECTS: usize = 3;
1870 assert_eq!(MAX_REDIRECTS, 3, "fetch_html allows exactly 3 redirects");
1871 }
1872
1873 #[test]
1876 fn redirect_no_location_error_message() {
1877 let err = std::io::Error::other("redirect with no Location");
1878 assert!(err.to_string().contains("redirect with no Location"));
1879 }
1880
1881 #[test]
1883 fn too_many_redirects_error_message() {
1884 let err = std::io::Error::other("too many redirects");
1885 assert!(err.to_string().contains("too many redirects"));
1886 }
1887
1888 #[test]
1890 fn non_2xx_status_error_format() {
1891 let status = reqwest::StatusCode::FORBIDDEN;
1892 let msg = format!("HTTP {status}");
1893 assert!(msg.contains("403"));
1894 }
1895
1896 #[test]
1898 fn not_found_status_error_format() {
1899 let status = reqwest::StatusCode::NOT_FOUND;
1900 let msg = format!("HTTP {status}");
1901 assert!(msg.contains("404"));
1902 }
1903
1904 #[test]
1906 fn relative_redirect_same_host_path() {
1907 let base = Url::parse("https://example.com/current").unwrap();
1908 let resolved = base.join("/other").unwrap();
1909 assert_eq!(resolved.as_str(), "https://example.com/other");
1910 }
1911
1912 #[test]
1914 fn relative_redirect_relative_path() {
1915 let base = Url::parse("https://example.com/a/b").unwrap();
1916 let resolved = base.join("c").unwrap();
1917 assert_eq!(resolved.as_str(), "https://example.com/a/c");
1918 }
1919
1920 #[test]
1922 fn absolute_redirect_overrides_base() {
1923 let base = Url::parse("https://example.com/page").unwrap();
1924 let resolved = base.join("https://other.com/target").unwrap();
1925 assert_eq!(resolved.as_str(), "https://other.com/target");
1926 }
1927
1928 #[test]
1930 fn redirect_http_downgrade_rejected() {
1931 let location = "http://example.com/page";
1932 let base = Url::parse("https://example.com/start").unwrap();
1933 let next = base.join(location).unwrap();
1934 let err = validate_url(next.as_str()).unwrap_err();
1935 assert_matches!(err, ToolError::Blocked { .. });
1936 }
1937
1938 #[test]
1940 fn redirect_location_private_ip_blocked() {
1941 let location = "https://192.168.100.1/admin";
1942 let base = Url::parse("https://example.com/start").unwrap();
1943 let next = base.join(location).unwrap();
1944 let err = validate_url(next.as_str()).unwrap_err();
1945 assert_matches!(err, ToolError::Blocked { .. });
1946 let ToolError::Blocked { command: cmd } = err else {
1947 panic!("expected Blocked");
1948 };
1949 assert!(
1950 cmd.contains("private") || cmd.contains("scheme"),
1951 "error message should describe the block reason: {cmd}"
1952 );
1953 }
1954
1955 #[test]
1957 fn redirect_location_internal_domain_blocked() {
1958 let location = "https://metadata.internal/latest/meta-data/";
1959 let base = Url::parse("https://example.com/start").unwrap();
1960 let next = base.join(location).unwrap();
1961 let err = validate_url(next.as_str()).unwrap_err();
1962 assert_matches!(err, ToolError::Blocked { .. });
1963 }
1964
1965 #[test]
1967 fn redirect_chain_three_hops_all_public() {
1968 let hops = [
1969 "https://redirect1.example.com/hop1",
1970 "https://redirect2.example.com/hop2",
1971 "https://destination.example.com/final",
1972 ];
1973 for hop in hops {
1974 assert!(validate_url(hop).is_ok(), "expected ok for {hop}");
1975 }
1976 }
1977
1978 #[test]
1983 fn redirect_to_private_ip_rejected_by_validate_url() {
1984 let private_targets = [
1986 "https://127.0.0.1/secret",
1987 "https://10.0.0.1/internal",
1988 "https://192.168.1.1/admin",
1989 "https://172.16.0.1/data",
1990 "https://[::1]/path",
1991 "https://[fe80::1]/path",
1992 "https://localhost/path",
1993 "https://service.internal/api",
1994 ];
1995 for target in private_targets {
1996 let result = validate_url(target);
1997 assert!(result.is_err(), "expected error for {target}");
1998 assert!(
1999 matches!(result.unwrap_err(), ToolError::Blocked { .. }),
2000 "expected Blocked for {target}"
2001 );
2002 }
2003 }
2004
2005 #[test]
2007 fn redirect_relative_url_resolves_correctly() {
2008 let base = Url::parse("https://example.com/page").unwrap();
2009 let relative = "/other";
2010 let resolved = base.join(relative).unwrap();
2011 assert_eq!(resolved.as_str(), "https://example.com/other");
2012 }
2013
2014 #[test]
2016 fn redirect_to_http_rejected() {
2017 let err = validate_url("http://example.com/page").unwrap_err();
2018 assert_matches!(err, ToolError::Blocked { .. });
2019 }
2020
2021 #[test]
2022 fn ipv4_mapped_ipv6_link_local_blocked() {
2023 let err = validate_url("https://[::ffff:169.254.0.1]/path").unwrap_err();
2024 assert_matches!(err, ToolError::Blocked { .. });
2025 }
2026
2027 #[test]
2028 fn ipv4_mapped_ipv6_public_allowed() {
2029 assert!(validate_url("https://[::ffff:93.184.216.34]/path").is_ok());
2030 }
2031
2032 #[tokio::test]
2035 async fn fetch_http_scheme_blocked() {
2036 let config = ScrapeConfig::default();
2037 let executor = WebScrapeExecutor::new(&config);
2038 let call = crate::executor::ToolCall {
2039 tool_id: ToolName::new("fetch"),
2040 params: {
2041 let mut m = serde_json::Map::new();
2042 m.insert("url".to_owned(), serde_json::json!("http://example.com"));
2043 m
2044 },
2045 caller_id: None,
2046 context: None,
2047
2048 tool_call_id: String::new(),
2049 skill_name: None,
2050 };
2051 let result = executor.execute_tool_call(&call).await;
2052 assert_matches!(result, Err(ToolError::Blocked { .. }));
2053 }
2054
2055 #[tokio::test]
2056 async fn fetch_private_ip_blocked() {
2057 let config = ScrapeConfig::default();
2058 let executor = WebScrapeExecutor::new(&config);
2059 let call = crate::executor::ToolCall {
2060 tool_id: ToolName::new("fetch"),
2061 params: {
2062 let mut m = serde_json::Map::new();
2063 m.insert(
2064 "url".to_owned(),
2065 serde_json::json!("https://192.168.1.1/secret"),
2066 );
2067 m
2068 },
2069 caller_id: None,
2070 context: None,
2071
2072 tool_call_id: String::new(),
2073 skill_name: None,
2074 };
2075 let result = executor.execute_tool_call(&call).await;
2076 assert_matches!(result, Err(ToolError::Blocked { .. }));
2077 }
2078
2079 #[tokio::test]
2080 async fn fetch_localhost_blocked() {
2081 let config = ScrapeConfig::default();
2082 let executor = WebScrapeExecutor::new(&config);
2083 let call = crate::executor::ToolCall {
2084 tool_id: ToolName::new("fetch"),
2085 params: {
2086 let mut m = serde_json::Map::new();
2087 m.insert(
2088 "url".to_owned(),
2089 serde_json::json!("https://localhost/page"),
2090 );
2091 m
2092 },
2093 caller_id: None,
2094 context: None,
2095
2096 tool_call_id: String::new(),
2097 skill_name: None,
2098 };
2099 let result = executor.execute_tool_call(&call).await;
2100 assert_matches!(result, Err(ToolError::Blocked { .. }));
2101 }
2102
2103 #[tokio::test]
2104 async fn fetch_unknown_tool_returns_none() {
2105 let config = ScrapeConfig::default();
2106 let executor = WebScrapeExecutor::new(&config);
2107 let call = crate::executor::ToolCall {
2108 tool_id: ToolName::new("unknown_tool"),
2109 params: serde_json::Map::new(),
2110 caller_id: None,
2111 context: None,
2112
2113 tool_call_id: String::new(),
2114 skill_name: None,
2115 };
2116 let result = executor.execute_tool_call(&call).await;
2117 assert!(result.unwrap().is_none());
2118 }
2119
2120 #[tokio::test]
2121 async fn fetch_returns_body_via_mock() {
2122 use wiremock::matchers::{method, path};
2123 use wiremock::{Mock, ResponseTemplate};
2124
2125 let (executor, server) = mock_server_executor().await;
2126 Mock::given(method("GET"))
2127 .and(path("/content"))
2128 .respond_with(ResponseTemplate::new(200).set_body_string("plain text content"))
2129 .mount(&server)
2130 .await;
2131
2132 let (host, addrs) = server_host_and_addr(&server);
2133 let url = format!("{}/content", server.uri());
2134 let result = executor
2135 .fetch_html(&url, &host, &addrs, "fetch", "test-cid", None, None)
2136 .await;
2137 assert!(result.is_ok());
2138 assert_eq!(result.unwrap(), "plain text content");
2139 }
2140
2141 #[test]
2142 fn tool_definitions_returns_web_scrape_and_fetch() {
2143 let config = ScrapeConfig::default();
2144 let executor = WebScrapeExecutor::new(&config);
2145 let defs = executor.tool_definitions();
2146 assert_eq!(defs.len(), 2);
2147 assert_eq!(defs[0].id, "web_scrape");
2148 assert_eq!(
2149 defs[0].invocation,
2150 crate::registry::InvocationHint::FencedBlock("scrape")
2151 );
2152 assert_eq!(defs[1].id, "fetch");
2153 assert_eq!(
2154 defs[1].invocation,
2155 crate::registry::InvocationHint::ToolCall
2156 );
2157 }
2158
2159 #[test]
2160 fn tool_definitions_schema_has_all_params() {
2161 let config = ScrapeConfig::default();
2162 let executor = WebScrapeExecutor::new(&config);
2163 let defs = executor.tool_definitions();
2164 let obj = defs[0].schema.as_object().unwrap();
2165 let props = obj["properties"].as_object().unwrap();
2166 assert!(props.contains_key("url"));
2167 assert!(props.contains_key("select"));
2168 assert!(props.contains_key("extract"));
2169 assert!(props.contains_key("limit"));
2170 let req = obj["required"].as_array().unwrap();
2171 assert!(req.iter().any(|v| v.as_str() == Some("url")));
2172 assert!(req.iter().any(|v| v.as_str() == Some("select")));
2173 assert!(!req.iter().any(|v| v.as_str() == Some("extract")));
2174 }
2175
2176 #[test]
2179 fn subdomain_localhost_blocked() {
2180 let host: url::Host<&str> = url::Host::Domain("foo.localhost");
2181 assert!(is_private_host(&host));
2182 }
2183
2184 #[test]
2185 fn internal_tld_blocked() {
2186 let host: url::Host<&str> = url::Host::Domain("service.internal");
2187 assert!(is_private_host(&host));
2188 }
2189
2190 #[test]
2191 fn local_tld_blocked() {
2192 let host: url::Host<&str> = url::Host::Domain("printer.local");
2193 assert!(is_private_host(&host));
2194 }
2195
2196 #[test]
2197 fn public_domain_not_blocked() {
2198 let host: url::Host<&str> = url::Host::Domain("example.com");
2199 assert!(!is_private_host(&host));
2200 }
2201
2202 #[tokio::test]
2205 async fn resolve_loopback_rejected() {
2206 let url = url::Url::parse("https://127.0.0.1/path").unwrap();
2208 let result = resolve_and_validate(&url).await;
2210 assert!(
2211 result.is_err(),
2212 "loopback IP must be rejected by resolve_and_validate"
2213 );
2214 let err = result.unwrap_err();
2215 assert_matches!(err, crate::executor::ToolError::Blocked { .. });
2216 }
2217
2218 #[tokio::test]
2219 async fn resolve_private_10_rejected() {
2220 let url = url::Url::parse("https://10.0.0.1/path").unwrap();
2221 let result = resolve_and_validate(&url).await;
2222 assert!(result.is_err());
2223 assert_matches!(
2224 result.unwrap_err(),
2225 crate::executor::ToolError::Blocked { .. }
2226 );
2227 }
2228
2229 #[tokio::test]
2230 async fn resolve_private_192_rejected() {
2231 let url = url::Url::parse("https://192.168.1.1/path").unwrap();
2232 let result = resolve_and_validate(&url).await;
2233 assert!(result.is_err());
2234 assert_matches!(
2235 result.unwrap_err(),
2236 crate::executor::ToolError::Blocked { .. }
2237 );
2238 }
2239
2240 #[tokio::test]
2241 async fn resolve_ipv6_loopback_rejected() {
2242 let url = url::Url::parse("https://[::1]/path").unwrap();
2243 let result = resolve_and_validate(&url).await;
2244 assert!(result.is_err());
2245 assert_matches!(
2246 result.unwrap_err(),
2247 crate::executor::ToolError::Blocked { .. }
2248 );
2249 }
2250
2251 #[tokio::test]
2252 async fn resolve_no_host_returns_ok() {
2253 let url = url::Url::parse("https://example.com/path").unwrap();
2255 let url_no_host = url::Url::parse("data:text/plain,hello").unwrap();
2257 let result = resolve_and_validate(&url_no_host).await;
2259 assert!(result.is_ok());
2260 let (host, addrs) = result.unwrap();
2261 assert!(host.is_empty());
2262 assert!(addrs.is_empty());
2263 drop(url);
2264 drop(url_no_host);
2265 }
2266
2267 async fn make_file_audit_logger(
2271 dir: &tempfile::TempDir,
2272 ) -> (
2273 std::sync::Arc<crate::audit::AuditLogger>,
2274 std::path::PathBuf,
2275 ) {
2276 use crate::audit::AuditLogger;
2277 use crate::config::AuditConfig;
2278 let path = dir.path().join("audit.log");
2279 let config = AuditConfig {
2280 enabled: true,
2281 destination: crate::config::AuditDestination::File(path.clone()),
2282 ..Default::default()
2283 };
2284 let logger = std::sync::Arc::new(AuditLogger::from_config(&config, false).await.unwrap());
2285 (logger, path)
2286 }
2287
2288 #[tokio::test]
2289 async fn with_audit_sets_logger() {
2290 let config = ScrapeConfig::default();
2291 let executor = WebScrapeExecutor::new(&config);
2292 assert!(executor.audit_logger.is_none());
2293
2294 let dir = tempfile::tempdir().unwrap();
2295 let (logger, _path) = make_file_audit_logger(&dir).await;
2296 let executor = executor.with_audit(logger);
2297 assert!(executor.audit_logger.is_some());
2298 }
2299
2300 #[test]
2301 fn tool_error_to_audit_result_blocked_maps_correctly() {
2302 let err = ToolError::Blocked {
2303 command: "scheme not allowed: http".into(),
2304 };
2305 let result = tool_error_to_audit_result(&err);
2306 assert!(
2307 matches!(result, AuditResult::Blocked { reason } if reason == "scheme not allowed: http")
2308 );
2309 }
2310
2311 #[test]
2312 fn tool_error_to_audit_result_timeout_maps_correctly() {
2313 let err = ToolError::Timeout { timeout_secs: 15 };
2314 let result = tool_error_to_audit_result(&err);
2315 assert_matches!(result, AuditResult::Timeout);
2316 }
2317
2318 #[test]
2319 fn tool_error_to_audit_result_execution_error_maps_correctly() {
2320 let err = ToolError::Execution(std::io::Error::other("connection refused"));
2321 let result = tool_error_to_audit_result(&err);
2322 assert!(
2323 matches!(result, AuditResult::Error { message } if message.contains("connection refused"))
2324 );
2325 }
2326
2327 #[tokio::test]
2328 async fn fetch_audit_blocked_url_logged() {
2329 let dir = tempfile::tempdir().unwrap();
2330 let (logger, log_path) = make_file_audit_logger(&dir).await;
2331
2332 let config = ScrapeConfig::default();
2333 let executor = WebScrapeExecutor::new(&config).with_audit(logger);
2334
2335 let call = crate::executor::ToolCall {
2336 tool_id: ToolName::new("fetch"),
2337 params: {
2338 let mut m = serde_json::Map::new();
2339 m.insert("url".to_owned(), serde_json::json!("http://example.com"));
2340 m
2341 },
2342 caller_id: None,
2343 context: None,
2344
2345 tool_call_id: String::new(),
2346 skill_name: None,
2347 };
2348 let result = executor.execute_tool_call(&call).await;
2349 assert_matches!(result, Err(ToolError::Blocked { .. }));
2350
2351 let content = tokio::fs::read_to_string(&log_path).await.unwrap();
2352 assert!(
2353 content.contains("\"tool\":\"fetch\""),
2354 "expected tool=fetch in audit: {content}"
2355 );
2356 assert!(
2357 content.contains("\"type\":\"blocked\""),
2358 "expected type=blocked in audit: {content}"
2359 );
2360 assert!(
2361 content.contains("http://example.com"),
2362 "expected URL in audit command field: {content}"
2363 );
2364 }
2365
2366 #[tokio::test]
2367 async fn log_audit_success_writes_to_file() {
2368 let dir = tempfile::tempdir().unwrap();
2369 let (logger, log_path) = make_file_audit_logger(&dir).await;
2370
2371 let config = ScrapeConfig::default();
2372 let executor = WebScrapeExecutor::new(&config).with_audit(logger);
2373
2374 executor
2375 .log_audit(
2376 "fetch",
2377 "https://example.com/page",
2378 AuditResult::Success,
2379 42,
2380 None,
2381 None,
2382 None,
2383 None,
2384 )
2385 .await;
2386
2387 let content = tokio::fs::read_to_string(&log_path).await.unwrap();
2388 assert!(
2389 content.contains("\"tool\":\"fetch\""),
2390 "expected tool=fetch in audit: {content}"
2391 );
2392 assert!(
2393 content.contains("\"type\":\"success\""),
2394 "expected type=success in audit: {content}"
2395 );
2396 assert!(
2397 content.contains("\"command\":\"https://example.com/page\""),
2398 "expected command URL in audit: {content}"
2399 );
2400 assert!(
2401 content.contains("\"duration_ms\":42"),
2402 "expected duration_ms in audit: {content}"
2403 );
2404 }
2405
2406 #[tokio::test]
2407 async fn log_audit_blocked_writes_to_file() {
2408 let dir = tempfile::tempdir().unwrap();
2409 let (logger, log_path) = make_file_audit_logger(&dir).await;
2410
2411 let config = ScrapeConfig::default();
2412 let executor = WebScrapeExecutor::new(&config).with_audit(logger);
2413
2414 executor
2415 .log_audit(
2416 "web_scrape",
2417 "http://evil.com/page",
2418 AuditResult::Blocked {
2419 reason: "scheme not allowed: http".into(),
2420 },
2421 0,
2422 None,
2423 None,
2424 None,
2425 None,
2426 )
2427 .await;
2428
2429 let content = tokio::fs::read_to_string(&log_path).await.unwrap();
2430 assert!(
2431 content.contains("\"tool\":\"web_scrape\""),
2432 "expected tool=web_scrape in audit: {content}"
2433 );
2434 assert!(
2435 content.contains("\"type\":\"blocked\""),
2436 "expected type=blocked in audit: {content}"
2437 );
2438 assert!(
2439 content.contains("scheme not allowed"),
2440 "expected block reason in audit: {content}"
2441 );
2442 }
2443
2444 #[tokio::test]
2445 async fn web_scrape_audit_blocked_url_logged() {
2446 let dir = tempfile::tempdir().unwrap();
2447 let (logger, log_path) = make_file_audit_logger(&dir).await;
2448
2449 let config = ScrapeConfig::default();
2450 let executor = WebScrapeExecutor::new(&config).with_audit(logger);
2451
2452 let call = crate::executor::ToolCall {
2453 tool_id: ToolName::new("web_scrape"),
2454 params: {
2455 let mut m = serde_json::Map::new();
2456 m.insert("url".to_owned(), serde_json::json!("http://example.com"));
2457 m.insert("select".to_owned(), serde_json::json!("h1"));
2458 m
2459 },
2460 caller_id: None,
2461 context: None,
2462
2463 tool_call_id: String::new(),
2464 skill_name: None,
2465 };
2466 let result = executor.execute_tool_call(&call).await;
2467 assert_matches!(result, Err(ToolError::Blocked { .. }));
2468
2469 let content = tokio::fs::read_to_string(&log_path).await.unwrap();
2470 assert!(
2471 content.contains("\"tool\":\"web_scrape\""),
2472 "expected tool=web_scrape in audit: {content}"
2473 );
2474 assert!(
2475 content.contains("\"type\":\"blocked\""),
2476 "expected type=blocked in audit: {content}"
2477 );
2478 }
2479
2480 #[tokio::test]
2481 async fn no_audit_logger_does_not_panic_on_blocked_fetch() {
2482 let config = ScrapeConfig::default();
2483 let executor = WebScrapeExecutor::new(&config);
2484 assert!(executor.audit_logger.is_none());
2485
2486 let call = crate::executor::ToolCall {
2487 tool_id: ToolName::new("fetch"),
2488 params: {
2489 let mut m = serde_json::Map::new();
2490 m.insert("url".to_owned(), serde_json::json!("http://example.com"));
2491 m
2492 },
2493 caller_id: None,
2494 context: None,
2495
2496 tool_call_id: String::new(),
2497 skill_name: None,
2498 };
2499 let result = executor.execute_tool_call(&call).await;
2501 assert_matches!(result, Err(ToolError::Blocked { .. }));
2502 }
2503
2504 #[tokio::test]
2506 async fn fetch_execute_tool_call_end_to_end() {
2507 use wiremock::matchers::{method, path};
2508 use wiremock::{Mock, ResponseTemplate};
2509
2510 let (executor, server) = mock_server_executor().await;
2511 Mock::given(method("GET"))
2512 .and(path("/e2e"))
2513 .respond_with(ResponseTemplate::new(200).set_body_string("<h1>end-to-end</h1>"))
2514 .mount(&server)
2515 .await;
2516
2517 let (host, addrs) = server_host_and_addr(&server);
2518 let result = executor
2520 .fetch_html(
2521 &format!("{}/e2e", server.uri()),
2522 &host,
2523 &addrs,
2524 "fetch",
2525 "test-cid",
2526 None,
2527 None,
2528 )
2529 .await;
2530 assert!(result.is_ok());
2531 assert!(result.unwrap().contains("end-to-end"));
2532 }
2533
2534 #[test]
2537 fn domain_matches_exact() {
2538 assert!(domain_matches("example.com", "example.com"));
2539 assert!(!domain_matches("example.com", "other.com"));
2540 assert!(!domain_matches("example.com", "sub.example.com"));
2541 }
2542
2543 #[test]
2544 fn domain_matches_wildcard_single_subdomain() {
2545 assert!(domain_matches("*.example.com", "sub.example.com"));
2546 assert!(!domain_matches("*.example.com", "example.com"));
2547 assert!(!domain_matches("*.example.com", "sub.sub.example.com"));
2548 }
2549
2550 #[test]
2551 fn domain_matches_wildcard_does_not_match_empty_label() {
2552 assert!(!domain_matches("*.example.com", ".example.com"));
2554 }
2555
2556 #[test]
2557 fn domain_matches_multi_wildcard_treated_as_exact() {
2558 assert!(!domain_matches("*.*.example.com", "a.b.example.com"));
2560 }
2561
2562 #[test]
2565 fn check_domain_policy_empty_lists_allow_all() {
2566 assert!(check_domain_policy("example.com", &[], &[]).is_ok());
2567 assert!(check_domain_policy("evil.com", &[], &[]).is_ok());
2568 }
2569
2570 #[test]
2571 fn check_domain_policy_denylist_blocks() {
2572 let denied = vec!["evil.com".to_string()];
2573 let err = check_domain_policy("evil.com", &[], &denied).unwrap_err();
2574 assert_matches!(err, ToolError::Blocked { .. });
2575 }
2576
2577 #[test]
2578 fn check_domain_policy_denylist_does_not_block_other_domains() {
2579 let denied = vec!["evil.com".to_string()];
2580 assert!(check_domain_policy("good.com", &[], &denied).is_ok());
2581 }
2582
2583 #[test]
2584 fn check_domain_policy_allowlist_permits_matching() {
2585 let allowed = vec!["docs.rs".to_string(), "*.rust-lang.org".to_string()];
2586 assert!(check_domain_policy("docs.rs", &allowed, &[]).is_ok());
2587 assert!(check_domain_policy("blog.rust-lang.org", &allowed, &[]).is_ok());
2588 }
2589
2590 #[test]
2591 fn check_domain_policy_allowlist_blocks_unknown() {
2592 let allowed = vec!["docs.rs".to_string()];
2593 let err = check_domain_policy("other.com", &allowed, &[]).unwrap_err();
2594 assert_matches!(err, ToolError::Blocked { .. });
2595 }
2596
2597 #[test]
2598 fn check_domain_policy_deny_overrides_allow() {
2599 let allowed = vec!["example.com".to_string()];
2600 let denied = vec!["example.com".to_string()];
2601 let err = check_domain_policy("example.com", &allowed, &denied).unwrap_err();
2602 assert_matches!(err, ToolError::Blocked { .. });
2603 }
2604
2605 #[test]
2606 fn check_domain_policy_wildcard_in_denylist() {
2607 let denied = vec!["*.evil.com".to_string()];
2608 let err = check_domain_policy("sub.evil.com", &[], &denied).unwrap_err();
2609 assert_matches!(err, ToolError::Blocked { .. });
2610 assert!(check_domain_policy("evil.com", &[], &denied).is_ok());
2612 }
2613}