lc_rag/loaders/
mod.rs

1// src/retrieval/loaders/mod.rs
2//! Document loader implementations
3//!
4//! Provides document loading from files in various formats, including PDF, CSV, Text, JSON,
5//! Markdown, HTML, etc. v0.4.1 added the WebScraper, Sitemap, and Docx loaders.
6
7mod csv;
8mod docx;
9mod html;
10mod json;
11mod markdown;
12mod pdf;
13mod sitemap;
14mod text;
15mod web_scraper;
16
17pub use csv::CSVLoader;
18pub use docx::DocxLoader;
19pub use html::HTMLLoader;
20pub use json::JSONLoader;
21pub use markdown::MarkdownLoader;
22pub use pdf::PDFLoader;
23pub use sitemap::SitemapLoader;
24pub use text::TextLoader;
25pub use web_scraper::WebScraperLoader;
26
27use async_trait::async_trait;
28use futures_util::{Stream, StreamExt};
29use lc_vector_stores::Document;
30use std::time::Duration;
31
32/// A5: cap the response body a URL-based loader will accept, so a hostile or
33/// misconfigured target cannot exhaust memory by streaming an unbounded body.
34const MAX_HTTP_BODY_BYTES: usize = 1024 * 1024; // 1 MiB
35
36/// A5: shared SSRF-hardened HTTP fetch used by the URL-based loaders
37/// (`HTMLLoader`, `WebScraperLoader`, `SitemapLoader`).
38///
39/// Redirects are *disabled* at the transport layer — manual redirect handling lives
40/// inside `lc_core::ssrf::guarded_get`, which resolves each hop once, validates every
41/// address, pins the validated IPs for the actual connection (DNS-rebinding closed),
42/// and applies the per-request timeout. The body is then read with a 1 MiB hard cap.
43/// Returns the post-redirect final URL and the body text so callers can record the
44/// true `source`.
45pub(crate) async fn guarded_fetch(
46    url: &str,
47    timeout: Duration,
48) -> Result<(String, String), LoaderError> {
49    let resp = lc_core::ssrf::guarded_get(url, true, Some(timeout))
50        .await
51        .map_err(|e| LoaderError::Other(format!("HTTP request failed {}: {}", url, e)))?;
52
53    let final_url = resp.url().as_str().to_string();
54
55    let status = resp.status();
56    if !status.is_success() {
57        return Err(LoaderError::Other(format!(
58            "HTTP error {}: {}",
59            url, status
60        )));
61    }
62
63    // Stream-capped read: enforce the 1 MiB limit incrementally instead of reading an
64    // unbounded body into memory first and only checking afterwards.
65    let body = read_capped_body(resp.bytes_stream(), url, MAX_HTTP_BODY_BYTES).await?;
66
67    let text = String::from_utf8(body)
68        .map_err(|_| LoaderError::Other(format!("response for {} is not valid UTF-8", url)))?;
69    Ok((final_url, text))
70}
71
72/// Drain a byte stream into a `Vec<u8>`, aborting as soon as the accumulated
73/// size exceeds `max_bytes`. The cap is checked *after every chunk*, so an
74/// oversized body is rejected before the remainder of the stream is polled —
75/// the caller (and the network peer) never have to buffer the whole response.
76/// (A5)
77pub(crate) async fn read_capped_body<S, T, E>(
78    mut stream: S,
79    url: &str,
80    max_bytes: usize,
81) -> Result<Vec<u8>, LoaderError>
82where
83    S: Stream<Item = Result<T, E>> + Unpin,
84    T: AsRef<[u8]>,
85    E: std::fmt::Display,
86{
87    let mut body = Vec::new();
88    while let Some(chunk) = stream.next().await {
89        let chunk = chunk
90            .map_err(|e| LoaderError::Other(format!("failed to read response {}: {}", url, e)))?;
91        body.extend_from_slice(chunk.as_ref());
92        if body.len() > max_bytes {
93            return Err(LoaderError::Other(format!(
94                "response for {} exceeds the {:.0} KiB size limit",
95                url,
96                max_bytes / 1024
97            )));
98        }
99    }
100    Ok(body)
101}
102
103/// Document loader error type
104#[derive(Debug, thiserror::Error)]
105#[non_exhaustive]
106pub enum LoaderError {
107    /// IO error
108    #[error("IO error: {0}")]
109    IoError(#[from] std::io::Error),
110
111    /// CSV parse error
112    #[error("CSV parse error: {0}")]
113    CsvError(String),
114
115    /// PDF parse error
116    #[error("PDF parse error: {0}")]
117    PdfError(String),
118
119    /// JSON parse error
120    #[error("JSON parse error: {0}")]
121    JsonError(String),
122
123    /// Unknown error
124    #[error("unknown error: {0}")]
125    Other(String),
126}
127
128impl From<pdf_extract::Error> for LoaderError {
129    fn from(err: pdf_extract::Error) -> Self {
130        LoaderError::PdfError(err.to_string())
131    }
132}
133
134/// Document loader trait
135///
136/// Defines the common interface for loading documents from a source.
137#[async_trait]
138pub trait DocumentLoader: Send + Sync {
139    /// Loads documents from the source
140    ///
141    /// # Returns
142    /// The loaded documents
143    async fn load(&self) -> Result<Vec<Document>, LoaderError>;
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    #[tokio::test]
151    async fn guarded_fetch_blocks_private_loopback() {
152        // A5: a loopback URL must be rejected by the SSRF guard before any request
153        // is sent — no network I/O happens, so this test needs no mock server.
154        let err = guarded_fetch("http://127.0.0.1/sitemap.xml", Duration::from_secs(5))
155            .await
156            .unwrap_err();
157        assert!(
158            err.to_string().contains("SSRF"),
159            "expected an SSRF rejection, got: {err}"
160        );
161    }
162
163    #[tokio::test]
164    async fn guarded_fetch_blocks_link_local() {
165        // A5: the cloud metadata address must be blocked even via IPv4-mapped IPv6.
166        let err = guarded_fetch(
167            "http://[::ffff:169.254.169.254]/latest",
168            Duration::from_secs(5),
169        )
170        .await
171        .unwrap_err();
172        assert!(
173            err.to_string().contains("SSRF"),
174            "expected an SSRF rejection, got: {err}"
175        );
176    }
177
178    #[tokio::test]
179    async fn read_capped_body_accepts_body_exactly_at_the_limit() {
180        use futures_util::stream;
181
182        let half = vec![b'a'; MAX_HTTP_BODY_BYTES / 2];
183        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> =
184            vec![Ok(half.clone()), Ok(half)];
185
186        let body = read_capped_body(
187            stream::iter(chunks),
188            "http://example.com/x",
189            MAX_HTTP_BODY_BYTES,
190        )
191        .await
192        .expect("exactly 1 MiB must be accepted");
193        assert_eq!(body.len(), MAX_HTTP_BODY_BYTES);
194    }
195
196    #[tokio::test]
197    async fn read_capped_body_rejects_oversize_without_polling_the_rest() {
198        use futures_util::stream;
199        use std::sync::atomic::{AtomicUsize, Ordering};
200        use std::sync::Arc;
201
202        // First chunk is already 600 KiB, the second pushes the total to 1.2 MiB.
203        // A third "poison" chunk records whether it was ever pulled from the
204        // stream — with incremental capping the reader must abort before that.
205        let poison_pulled = Arc::new(AtomicUsize::new(0));
206        let counter = poison_pulled.clone();
207        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> = vec![
208            Ok(vec![b'a'; 600 * 1024]),
209            Ok(vec![b'b'; 600 * 1024]),
210            Ok(vec![b'c'; 1024]),
211        ];
212        let s = stream::unfold(0usize, move |idx| {
213            let chunks = chunks.clone();
214            let counter = counter.clone();
215            async move {
216                if idx >= chunks.len() {
217                    return None;
218                }
219                if idx == 2 {
220                    counter.fetch_add(1, Ordering::SeqCst);
221                }
222                let item = chunks[idx].clone();
223                Some((item, idx + 1))
224            }
225        });
226        // unfold streams are not `Unpin` (they pin the in-flight future); reqwest's
227        // byte stream is Unpin, and boxing matches the production bound.
228        let s = Box::pin(s);
229
230        let err = read_capped_body(s, "http://example.com/huge", MAX_HTTP_BODY_BYTES)
231            .await
232            .unwrap_err();
233        assert!(
234            err.to_string().contains("1024 KiB size limit"),
235            "expected a size-limit error, got: {err}"
236        );
237        assert_eq!(
238            poison_pulled.load(Ordering::SeqCst),
239            0,
240            "stream must not be polled for chunks past the cap (unbounded buffering)"
241        );
242    }
243
244    #[tokio::test]
245    async fn read_capped_body_rejects_non_utf8_inside_the_limit() {
246        use futures_util::stream;
247
248        let chunks: Vec<Result<Vec<u8>, std::convert::Infallible>> =
249            vec![Ok(b"abc".to_vec()), Ok(vec![0xff, 0xfe])];
250
251        // The capped reader itself returns raw bytes; mirror guarded_fetch's
252        // UTF-8 conversion to lock the surfaced error.
253        let body = read_capped_body(stream::iter(chunks), "http://example.com/bin", 64)
254            .await
255            .expect("small body accepted");
256        let err = String::from_utf8(body).unwrap_err();
257        assert!(err.to_string().contains("invalid utf-8"));
258    }
259}
lc_rag/loaders/mod.rs

lc_rag/loaders/
mod.rs
