olai_http/
lib.rs

1#[cfg(feature = "recording")]
2use std::collections::HashMap;
3#[cfg(feature = "recording")]
4use std::path::PathBuf;
5use std::sync::Arc;
6#[cfg(feature = "recording")]
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::Duration;
9
10use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
11use reqwest::{Body, Client, IntoUrl, Method, RequestBuilder};
12use serde::Serialize;
13use tokio::runtime::Handle;
14
15use self::service::{HttpService, make_service};
16pub use self::token::{TemporaryToken, TokenCache};
17
18pub mod aws;
19pub mod azure;
20mod backoff;
21mod client;
22mod config;
23mod credential;
24pub mod databricks;
25mod error;
26pub mod gcp;
27pub mod service;
28
29mod retry;
30mod token;
31mod util;
32
33pub use client::{Certificate, ClientConfigKey, ClientOptions};
34pub use credential::*;
35pub use error::*;
36pub use retry::RetryConfig;
37pub use service::{ReqwestService, SpawnService};
38
39/// A shared, cloneable signer for a `CloudClient`.
40type SharedSigner = Arc<dyn RequestSigner>;
41
42/// A no-op signer used by `new_unauthenticated`.
43#[derive(Debug)]
44struct NoopSigner;
45
46impl RequestSigner for NoopSigner {
47    fn sign<'a>(
48        &'a self,
49        req: RequestBuilder,
50    ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
51        Box::pin(async move { Ok(req) })
52    }
53}
54
55/// A signer that injects a static bearer token.
56#[derive(Debug)]
57struct BearerTokenSigner {
58    token: String,
59}
60
61impl RequestSigner for BearerTokenSigner {
62    fn sign<'a>(
63        &'a self,
64        req: RequestBuilder,
65    ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
66        let token = self.token.clone();
67        Box::pin(async move { Ok(req.bearer_auth(&token)) })
68    }
69}
70
71#[cfg(feature = "recording")]
72#[derive(Debug, Clone)]
73struct RecordingState {
74    out_dir: PathBuf,
75    counter: Arc<AtomicU64>,
76}
77
78/// An authenticated HTTP client for cloud provider APIs.
79///
80/// Created via the provider-specific constructors [`CloudClient::new_aws`],
81/// [`CloudClient::new_azure`], [`CloudClient::new_google`], or
82/// [`CloudClient::new_databricks`], or the simpler [`CloudClient::new_with_token`]
83/// and [`CloudClient::new_unauthenticated`].
84#[derive(Clone)]
85pub struct CloudClient {
86    signer: SharedSigner,
87    reqwest_client: Client,
88    service: Arc<dyn HttpService>,
89    /// Retry configuration stored for future use by [`CloudClient::send`].
90    ///
91    /// Currently the retry policy is applied inside credential providers (token
92    /// refresh), but is not yet applied to user-initiated requests. Stored here
93    /// so that a `with_retry_config()` builder method can expose it without a
94    /// breaking change.
95    pub retry_config: RetryConfig,
96    #[cfg(feature = "recording")]
97    recording: Option<RecordingState>,
98}
99
100impl CloudClient {
101    fn new_with_signer(
102        signer: SharedSigner,
103        reqwest_client: Client,
104        service: Arc<dyn HttpService>,
105        retry_config: RetryConfig,
106    ) -> Self {
107        Self {
108            signer,
109            reqwest_client,
110            service,
111            retry_config,
112            #[cfg(feature = "recording")]
113            recording: None,
114        }
115    }
116
117    /// Create a new client with AWS credentials.
118    ///
119    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
120    /// will be spawned on the given runtime handle.
121    pub fn new_aws<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
122    where
123        I: IntoIterator<Item = (K, V)>,
124        K: AsRef<str>,
125        V: Into<String>,
126    {
127        let config = options
128            .into_iter()
129            .fold(
130                aws::AmazonBuilder::new(),
131                |builder, (key, value)| match key.as_ref().parse() {
132                    Ok(k) => builder.with_config(k, value),
133                    Err(_) => builder,
134                },
135            )
136            .build(runtime)?;
137
138        let reqwest_client = config.client_options.client()?;
139        let service = make_service(reqwest_client.clone(), runtime);
140        let retry_config = config.retry_config.clone();
141        Ok(Self::new_with_signer(
142            Arc::new(config),
143            reqwest_client,
144            service,
145            retry_config,
146        ))
147    }
148
149    /// Create a new client with Google Cloud credentials.
150    ///
151    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
152    /// will be spawned on the given runtime handle.
153    pub fn new_google<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
154    where
155        I: IntoIterator<Item = (K, V)>,
156        K: AsRef<str>,
157        V: Into<String>,
158    {
159        let config = options
160            .into_iter()
161            .fold(
162                gcp::GoogleBuilder::new(),
163                |builder, (key, value)| match key.as_ref().parse() {
164                    Ok(k) => builder.with_config(k, value),
165                    Err(_) => builder,
166                },
167            )
168            .build(runtime)?;
169
170        let reqwest_client = config.client_options.client()?;
171        let service = make_service(reqwest_client.clone(), runtime);
172        let retry_config = config.retry_config.clone();
173        Ok(Self::new_with_signer(
174            Arc::new(config),
175            reqwest_client,
176            service,
177            retry_config,
178        ))
179    }
180
181    /// Create a new client with Azure credentials.
182    ///
183    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
184    /// will be spawned on the given runtime handle.
185    pub fn new_azure<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
186    where
187        I: IntoIterator<Item = (K, V)>,
188        K: AsRef<str>,
189        V: Into<String>,
190    {
191        let config = options
192            .into_iter()
193            .fold(
194                azure::AzureBuilder::new(),
195                |builder, (key, value)| match key.as_ref().parse() {
196                    Ok(k) => builder.with_config(k, value),
197                    Err(_) => builder,
198                },
199            )
200            .build(runtime)?;
201
202        let reqwest_client = config.client_options.client()?;
203        let service = make_service(reqwest_client.clone(), runtime);
204        let retry_config = config.retry_config.clone();
205        Ok(Self::new_with_signer(
206            Arc::new(config),
207            reqwest_client,
208            service,
209            retry_config,
210        ))
211    }
212
213    /// Create a new client with Databricks credentials.
214    ///
215    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
216    /// will be spawned on the given runtime handle.
217    pub fn new_databricks<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
218    where
219        I: IntoIterator<Item = (K, V)>,
220        K: AsRef<str>,
221        V: Into<String>,
222    {
223        use databricks::DatabricksBuilder;
224
225        let config = options
226            .into_iter()
227            .fold(
228                DatabricksBuilder::new(),
229                |builder, (key, value)| match key.as_ref().parse() {
230                    Ok(k) => builder.with_config(k, value),
231                    Err(_) => builder,
232                },
233            )
234            .build(runtime)?;
235
236        let reqwest_client = config.client_options.client()?;
237        let service = make_service(reqwest_client.clone(), runtime);
238        let retry_config = config.retry_config.clone();
239        Ok(Self::new_with_signer(
240            Arc::new(config),
241            reqwest_client,
242            service,
243            retry_config,
244        ))
245    }
246
247    /// Create a new client with a personal access token.
248    pub fn new_with_token(token: impl ToString) -> Self {
249        let reqwest_client = Client::new();
250        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
251        Self::new_with_signer(
252            Arc::new(BearerTokenSigner {
253                token: token.to_string(),
254            }),
255            reqwest_client,
256            service,
257            RetryConfig::default(),
258        )
259    }
260
261    /// Create a new unauthenticated client.
262    pub fn new_unauthenticated() -> Self {
263        let reqwest_client = Client::new();
264        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
265        Self::new_with_signer(
266            Arc::new(NoopSigner),
267            reqwest_client,
268            service,
269            RetryConfig::default(),
270        )
271    }
272
273    /// Route all HTTP I/O through the given runtime handle.
274    ///
275    /// This is useful for simple constructors (`new_with_token`, `new_unauthenticated`)
276    /// where no credential providers need to perform HTTP I/O. For cloud provider
277    /// constructors, pass the handle at construction time instead.
278    pub fn with_runtime(mut self, handle: Handle) -> Self {
279        self.service = Arc::new(SpawnService::new(self.service, handle));
280        self
281    }
282
283    /// Replace the [`HttpService`] used for request execution.
284    pub fn with_http_service(mut self, service: Arc<dyn HttpService>) -> Self {
285        self.service = service;
286        self
287    }
288
289    pub fn request<U: IntoUrl>(&self, method: Method, url: U) -> CloudRequestBuilder {
290        CloudRequestBuilder {
291            builder: self.reqwest_client.request(method, url),
292            client: self.clone(),
293            #[cfg(feature = "recording")]
294            out_dir: self.recording.as_ref().map(|r| r.out_dir.clone()),
295        }
296    }
297
298    pub fn get<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
299        self.request(Method::GET, url)
300    }
301
302    pub fn post<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
303        self.request(Method::POST, url)
304    }
305
306    pub fn put<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
307        self.request(Method::PUT, url)
308    }
309
310    pub fn delete<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
311        self.request(Method::DELETE, url)
312    }
313
314    pub fn head<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
315        self.request(Method::HEAD, url)
316    }
317
318    pub fn patch<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
319        self.request(Method::PATCH, url)
320    }
321
322    pub fn options<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
323        self.request(Method::OPTIONS, url)
324    }
325
326    pub fn trace<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
327        self.request(Method::TRACE, url)
328    }
329
330    pub fn connect<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
331        self.request(Method::CONNECT, url)
332    }
333
334    #[cfg(feature = "recording")]
335    pub fn set_recording_dir(&mut self, out_dir: std::path::PathBuf) -> Result<(), std::io::Error> {
336        let out_dir = std::fs::canonicalize(out_dir)?;
337        self.recording = Some(RecordingState {
338            out_dir,
339            counter: Arc::new(AtomicU64::new(0)),
340        });
341        Ok(())
342    }
343}
344
345pub struct CloudRequestBuilder {
346    builder: RequestBuilder,
347    client: CloudClient,
348    #[cfg(feature = "recording")]
349    out_dir: Option<PathBuf>,
350}
351
352impl CloudRequestBuilder {
353    /// Add a `Header` to this Request.
354    pub fn header<K, V>(mut self, key: K, value: V) -> CloudRequestBuilder
355    where
356        HeaderName: TryFrom<K>,
357        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
358        HeaderValue: TryFrom<V>,
359        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
360    {
361        self.builder = self.builder.header(key, value);
362        self
363    }
364
365    /// Add a set of Headers to the existing ones on this Request.
366    ///
367    /// The headers will be merged in to any already set.
368    pub fn headers(mut self, headers: HeaderMap) -> CloudRequestBuilder {
369        self.builder = self.builder.headers(headers);
370        self
371    }
372
373    /// Set the request body.
374    pub fn body<T: Into<Body>>(mut self, body: T) -> CloudRequestBuilder {
375        self.builder = self.builder.body(body);
376        self
377    }
378
379    /// Enables a request timeout.
380    ///
381    /// The timeout is applied from when the request starts connecting until the
382    /// response body has finished. It affects only this request and overrides
383    /// the timeout configured using `ClientBuilder::timeout()`.
384    pub fn timeout(mut self, timeout: Duration) -> CloudRequestBuilder {
385        self.builder = self.builder.timeout(timeout);
386        self
387    }
388
389    /// Modify the query string of the URL.
390    ///
391    /// Modifies the URL of this request, adding the parameters provided.
392    /// This method appends and does not overwrite. This means that it can
393    /// be called multiple times and that existing query parameters are not
394    /// overwritten if the same key is used. The key will simply show up
395    /// twice in the query string.
396    /// Calling `.query(&[("foo", "a"), ("foo", "b")])` gives `"foo=a&foo=b"`.
397    ///
398    /// # Note
399    /// This method does not support serializing a single key-value
400    /// pair. Instead of using `.query(("key", "val"))`, use a sequence, such
401    /// as `.query(&[("key", "val")])`. It's also possible to serialize structs
402    /// and maps into a key-value pair.
403    ///
404    /// # Errors
405    /// This method will fail if the object you provide cannot be serialized
406    /// into a query string.
407    pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> CloudRequestBuilder {
408        self.builder = self.builder.query(query);
409        self
410    }
411
412    /// Send a JSON body.
413    ///
414    /// # Errors
415    ///
416    /// Serialization can fail if `T`'s implementation of `Serialize` decides to
417    /// fail, or if `T` contains a map with non-string keys.
418    pub fn json<T: Serialize + ?Sized>(mut self, json: &T) -> CloudRequestBuilder {
419        self.builder = self.builder.json(json);
420        self
421    }
422
423    pub async fn send(mut self) -> Result<reqwest::Response> {
424        self.builder = self.client.signer.sign(self.builder).await?;
425
426        #[cfg(not(feature = "recording"))]
427        {
428            let request = self.builder.build()?;
429            let response = self.client.service.call(request).await?;
430            Ok(response)
431        }
432        #[cfg(feature = "recording")]
433        {
434            let response = send_record(self).await?;
435            Ok(response)
436        }
437    }
438}
439
440#[cfg(feature = "recording")]
441#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
442pub struct RequestResponseInfo {
443    pub request: RequestInfo,
444    pub response: ResponseInfo,
445}
446
447#[cfg(feature = "recording")]
448#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
449pub struct RequestInfo {
450    pub method: String,
451    pub url_path: String,
452    pub body: Option<String>,
453}
454
455#[cfg(feature = "recording")]
456#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
457pub struct ResponseInfo {
458    pub status: u16,
459    pub headers: HashMap<String, String>,
460    pub body: Option<String>,
461}
462
463/// Header names whose values must never appear in recording files.
464///
465/// These headers carry bearer tokens, signing secrets, or session tokens.
466/// They are replaced with `"<REDACTED>"` before any recording is written to disk.
467#[cfg(feature = "recording")]
468const SENSITIVE_HEADERS: &[&str] = &[
469    "authorization",
470    "x-amz-security-token",
471    "x-amz-content-sha256",
472    "x-databricks-authorization",
473    "x-ms-identity-principal-id",
474    "x-goog-iam-credentials-token",
475    "cookie",
476    "set-cookie",
477];
478
479#[cfg(feature = "recording")]
480fn redact_headers(headers: &HashMap<String, String>) -> HashMap<String, String> {
481    headers
482        .iter()
483        .map(|(k, v)| {
484            let v = if SENSITIVE_HEADERS.contains(&k.to_lowercase().as_str()) {
485                "<REDACTED>".to_string()
486            } else {
487                v.clone()
488            };
489            (k.clone(), v)
490        })
491        .collect()
492}
493
494#[cfg(feature = "recording")]
495async fn send_record(builder: CloudRequestBuilder) -> Result<reqwest::Response, reqwest::Error> {
496    let Some(out_dir) = builder.out_dir else {
497        let request = builder.builder.build().expect("request to be valid");
498        return builder.client.service.call(request).await;
499    };
500    let (_client, request) = builder.builder.build_split();
501    let request = request.expect("request to be valid");
502
503    let request_info = RequestInfo {
504        method: request.method().as_str().to_string(),
505        url_path: {
506            let url = request.url();
507            match url.query() {
508                Some(query) => format!("{}?{}", url.path(), query),
509                None => url.path().to_string(),
510            }
511        },
512        body: request
513            .body()
514            .and_then(|b| b.as_bytes().map(|b| String::from_utf8_lossy(b).to_string())),
515    };
516
517    let response = builder.client.service.call(request).await?;
518
519    // Record the response
520    let status = response.status().as_u16();
521    let raw_headers: HashMap<String, String> = response
522        .headers()
523        .iter()
524        .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
525        .collect();
526
527    // Get response body while preserving it for the caller
528    let response_bytes = response.bytes().await?;
529    let response_body = if response_bytes.is_empty() {
530        None
531    } else {
532        Some(String::from_utf8_lossy(&response_bytes).to_string())
533    };
534
535    let recording = RequestResponseInfo {
536        request: request_info,
537        response: ResponseInfo {
538            status,
539            // Redact sensitive headers before writing to disk
540            headers: redact_headers(&raw_headers),
541            body: response_body,
542        },
543    };
544
545    let counter = builder
546        .client
547        .recording
548        .as_ref()
549        .map(|r| r.counter.fetch_add(1, Ordering::SeqCst))
550        .unwrap_or(0);
551    let file_path = out_dir.join(format!("{:04}.json", counter));
552    if let Err(e) = std::fs::File::create(&file_path)
553        .and_then(|f| serde_json::to_writer_pretty(f, &recording).map_err(Into::into))
554    {
555        tracing::warn!(
556            "Failed to write recording to {}: {}",
557            file_path.display(),
558            e
559        );
560    }
561
562    // Return a new response built from the recorded data, using the raw (unredacted) headers
563    let mut mock_response = http::Response::builder().status(status);
564    for (k, v) in &raw_headers {
565        mock_response = mock_response.header(k, v);
566    }
567    let mock_response = mock_response
568        .body(response_bytes)
569        .expect("valid status code and headers");
570
571    Ok(reqwest::Response::from(mock_response))
572}
573
574#[cfg(all(test, feature = "recording"))]
575mod tests {
576    use super::*;
577    use std::fs;
578    use tempfile::TempDir;
579
580    #[tokio::test]
581    async fn test_request_response_recording() {
582        // Create a temporary directory for recordings
583        let temp_dir = TempDir::new().unwrap();
584        let temp_path = temp_dir.path().to_path_buf();
585
586        // Set up a mock server
587        let mut server = mockito::Server::new_async().await;
588        let mock = server
589            .mock("GET", "/test")
590            .with_status(200)
591            .with_header("content-type", "application/json")
592            .with_body(r#"{"message": "Hello, World!"}"#)
593            .create_async()
594            .await;
595
596        // Create a cloud client with recording enabled
597        let mut client = CloudClient::new_unauthenticated();
598        client.set_recording_dir(temp_path.clone()).unwrap();
599
600        // Make a request
601        let url = format!("{}/test", server.url());
602        let response = client.get(&url).send().await.unwrap();
603
604        // Verify the response is correct
605        assert_eq!(response.status(), 200);
606        let body = response.text().await.unwrap();
607        assert_eq!(body, r#"{"message": "Hello, World!"}"#);
608
609        // Verify that a recording file was created
610        let recordings: Vec<_> = fs::read_dir(&temp_path)
611            .unwrap()
612            .filter_map(|entry| {
613                let entry = entry.ok()?;
614                let path = entry.path();
615                if path.extension()? == "json" {
616                    Some(path)
617                } else {
618                    None
619                }
620            })
621            .collect();
622
623        assert_eq!(recordings.len(), 1, "Expected exactly one recording file");
624
625        // Read and verify the recording content
626        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
627        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
628
629        // Verify request information
630        assert_eq!(recording.request.method, "GET");
631        assert_eq!(recording.request.url_path, "/test");
632        assert_eq!(recording.request.body, None);
633
634        // Verify response information
635        assert_eq!(recording.response.status, 200);
636        assert_eq!(
637            recording.response.headers.get("content-type").unwrap(),
638            "application/json"
639        );
640        assert_eq!(
641            recording.response.body.as_ref().unwrap(),
642            r#"{"message": "Hello, World!"}"#
643        );
644
645        mock.assert_async().await;
646    }
647
648    #[tokio::test]
649    async fn test_recording_with_request_body() {
650        // Create a temporary directory for recordings
651        let temp_dir = TempDir::new().unwrap();
652        let temp_path = temp_dir.path().to_path_buf();
653
654        // Set up a mock server
655        let mut server = mockito::Server::new_async().await;
656        let mock = server
657            .mock("POST", "/create")
658            .with_status(201)
659            .with_header("location", "/resource/123")
660            .with_body(r#"{"id": 123, "status": "created"}"#)
661            .create_async()
662            .await;
663
664        // Create a cloud client with recording enabled
665        let mut client = CloudClient::new_unauthenticated();
666        client.set_recording_dir(temp_path.clone()).unwrap();
667
668        // Make a POST request with body
669        let url = format!("{}/create", server.url());
670        let response = client
671            .post(&url)
672            .json(&serde_json::json!({"name": "test resource"}))
673            .send()
674            .await
675            .unwrap();
676
677        // Verify the response
678        assert_eq!(response.status(), 201);
679
680        // Verify that a recording file was created
681        let recordings: Vec<_> = fs::read_dir(&temp_path)
682            .unwrap()
683            .filter_map(|entry| {
684                let entry = entry.ok()?;
685                let path = entry.path();
686                if path.extension()? == "json" {
687                    Some(path)
688                } else {
689                    None
690                }
691            })
692            .collect();
693
694        assert_eq!(recordings.len(), 1);
695
696        // Read and verify the recording content
697        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
698        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
699
700        // Verify request information
701        assert_eq!(recording.request.method, "POST");
702        assert_eq!(recording.request.url_path, "/create");
703        assert!(recording.request.body.is_some());
704        assert!(recording.request.body.unwrap().contains("test resource"));
705
706        // Verify response information
707        assert_eq!(recording.response.status, 201);
708        assert_eq!(
709            recording.response.headers.get("location").unwrap(),
710            "/resource/123"
711        );
712        assert!(recording.response.body.unwrap().contains("created"));
713
714        mock.assert_async().await;
715    }
716
717    #[tokio::test]
718    async fn test_counter_based_file_naming() {
719        // Create a temporary directory for recordings
720        let temp_dir = TempDir::new().unwrap();
721        let temp_path = temp_dir.path().to_path_buf();
722
723        // Set up a mock server
724        let mut server = mockito::Server::new_async().await;
725        let mock1 = server
726            .mock("GET", "/first")
727            .with_status(200)
728            .with_body("first response")
729            .create_async()
730            .await;
731        let mock2 = server
732            .mock("GET", "/second")
733            .with_status(200)
734            .with_body("second response")
735            .create_async()
736            .await;
737
738        // Create a cloud client with recording enabled
739        let mut client = CloudClient::new_unauthenticated();
740        client.set_recording_dir(temp_path.clone()).unwrap();
741
742        // Make multiple requests
743        let url1 = format!("{}/first", server.url());
744        let url2 = format!("{}/second", server.url());
745
746        let _response1 = client.get(&url1).send().await.unwrap();
747        let _response2 = client.get(&url2).send().await.unwrap();
748
749        // Verify that files are named with incrementing counter
750        let mut recordings: Vec<_> = fs::read_dir(&temp_path)
751            .unwrap()
752            .filter_map(|entry| {
753                let entry = entry.ok()?;
754                let path = entry.path();
755                if path.extension()? == "json" {
756                    Some(path)
757                } else {
758                    None
759                }
760            })
761            .collect();
762
763        recordings.sort();
764        assert_eq!(recordings.len(), 2);
765
766        // Check that files are named 000000.json and 000001.json
767        assert!(recordings[0].file_name().unwrap().to_str().unwrap() == "0000.json");
768        assert!(recordings[1].file_name().unwrap().to_str().unwrap() == "0001.json");
769
770        // Verify content matches the order of requests
771        let first_content = fs::read_to_string(&recordings[0]).unwrap();
772        let first_recording: RequestResponseInfo = serde_json::from_str(&first_content).unwrap();
773        assert_eq!(first_recording.request.url_path, "/first");
774        assert_eq!(
775            first_recording.response.body.as_ref().unwrap(),
776            "first response"
777        );
778
779        let second_content = fs::read_to_string(&recordings[1]).unwrap();
780        let second_recording: RequestResponseInfo = serde_json::from_str(&second_content).unwrap();
781        assert_eq!(second_recording.request.url_path, "/second");
782        assert_eq!(
783            second_recording.response.body.as_ref().unwrap(),
784            "second response"
785        );
786
787        mock1.assert_async().await;
788        mock2.assert_async().await;
789    }
790
791    #[tokio::test]
792    async fn test_query_parameter_recording() {
793        // Create a temporary directory for recordings
794        let temp_dir = TempDir::new().unwrap();
795        let temp_path = temp_dir.path().to_path_buf();
796
797        // Start a mock server
798        let mut server = mockito::Server::new_async().await;
799
800        // Create a mock that expects query parameters
801        let mock = server
802            .mock("GET", "/catalogs?max_results=10&page_token=abc123")
803            .with_status(200)
804            .with_header("content-type", "application/json")
805            .with_body(r#"{"catalogs": []}"#)
806            .create_async()
807            .await;
808
809        // Create a client with recording enabled
810        let mut client = CloudClient::new_unauthenticated();
811        client.set_recording_dir(temp_path.clone()).unwrap();
812
813        // Make a request with query parameters
814        let url = format!("{}/catalogs?max_results=10&page_token=abc123", server.url());
815        let response = client.get(&url).send().await.unwrap();
816
817        assert!(response.status().is_success());
818
819        // Verify that the recording file was created
820        let recordings: Vec<_> = fs::read_dir(&temp_path)
821            .unwrap()
822            .filter_map(|entry| {
823                let entry = entry.ok()?;
824                let path = entry.path();
825                if path.extension()? == "json" {
826                    Some(path)
827                } else {
828                    None
829                }
830            })
831            .collect();
832
833        assert_eq!(recordings.len(), 1);
834
835        // Read and verify the recording content includes query parameters
836        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
837        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
838
839        // Verify request information includes query parameters
840        assert_eq!(recording.request.method, "GET");
841        assert_eq!(
842            recording.request.url_path,
843            "/catalogs?max_results=10&page_token=abc123"
844        );
845        assert_eq!(recording.request.body, None);
846
847        // Verify response information
848        assert_eq!(recording.response.status, 200);
849        assert_eq!(
850            recording.response.body.as_ref().unwrap(),
851            r#"{"catalogs": []}"#
852        );
853
854        mock.assert_async().await;
855    }
856
857    /// Verify that the bearer token injected by a signed request does not appear
858    /// in the recording file. Request headers are currently not recorded, so this
859    /// test confirms the token does not leak via any other path (e.g. echoed back
860    /// in a response header or response body).
861    #[tokio::test]
862    async fn test_recording_does_not_contain_bearer_token_value() {
863        let temp_dir = TempDir::new().unwrap();
864        let temp_path = temp_dir.path().to_path_buf();
865
866        let mut server = mockito::Server::new_async().await;
867        let mock = server
868            .mock("GET", "/secret")
869            .match_header("authorization", mockito::Matcher::Any)
870            .with_status(200)
871            // Server does NOT echo the token back — simulates a well-behaved API
872            .with_body("ok")
873            .create_async()
874            .await;
875
876        let mut client = CloudClient::new_with_token("super-secret-token-12345");
877        client.set_recording_dir(temp_path.clone()).unwrap();
878
879        let url = format!("{}/secret", server.url());
880        client.get(&url).send().await.unwrap();
881
882        let recording_path = temp_path.join("0000.json");
883        let content = fs::read_to_string(&recording_path).unwrap();
884
885        // The raw token must not appear in the file at all
886        assert!(
887            !content.contains("super-secret-token-12345"),
888            "raw bearer token leaked into recording: {content}"
889        );
890
891        mock.assert_async().await;
892    }
893
894    /// Verify that sensitive headers returned by the server (e.g. a reflected
895    /// Authorization or AWS security token) are redacted before being written
896    /// to the recording file.
897    #[tokio::test]
898    async fn test_recording_redacts_sensitive_response_headers() {
899        let temp_dir = TempDir::new().unwrap();
900        let temp_path = temp_dir.path().to_path_buf();
901
902        let mut server = mockito::Server::new_async().await;
903        // Simulate a server that echoes a sensitive header back in its response
904        let mock = server
905            .mock("GET", "/s3")
906            .with_status(200)
907            .with_header("x-amz-security-token", "AQoXnyc4LLI2AJvUAMOGAR8a1234567890")
908            .with_header("authorization", "Bearer should-be-redacted")
909            .with_body("{}")
910            .create_async()
911            .await;
912
913        let mut client = CloudClient::new_unauthenticated();
914        client.set_recording_dir(temp_path.clone()).unwrap();
915
916        let url = format!("{}/s3", server.url());
917        client.get(&url).send().await.unwrap();
918
919        let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
920        assert!(
921            !content.contains("AQoXnyc4LLI2AJvUAMOGAR8a1234567890"),
922            "x-amz-security-token leaked into recording: {content}"
923        );
924        assert!(
925            !content.contains("should-be-redacted"),
926            "Authorization value leaked into recording: {content}"
927        );
928        // Both headers should be replaced with the redaction sentinel
929        assert_eq!(
930            content.matches("<REDACTED>").count(),
931            2,
932            "expected 2 <REDACTED> entries in recording: {content}"
933        );
934
935        mock.assert_async().await;
936    }
937
938    /// Verify that a recording produced with redacted headers can still be parsed.
939    #[tokio::test]
940    async fn test_recording_remains_valid_json_after_redaction() {
941        let temp_dir = TempDir::new().unwrap();
942        let temp_path = temp_dir.path().to_path_buf();
943
944        let mut server = mockito::Server::new_async().await;
945        let _mock = server
946            .mock("GET", "/check")
947            .with_status(200)
948            .with_header("authorization", "Bearer should-be-redacted")
949            .with_body(r#"{"ok":true}"#)
950            .create_async()
951            .await;
952
953        let mut client = CloudClient::new_unauthenticated();
954        client.set_recording_dir(temp_path.clone()).unwrap();
955
956        let url = format!("{}/check", server.url());
957        client.get(&url).send().await.unwrap();
958
959        let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
960        // Must parse cleanly as a RequestResponseInfo
961        let parsed: RequestResponseInfo = serde_json::from_str(&content)
962            .expect("recording file must be valid JSON even after redaction");
963        assert_eq!(parsed.response.status, 200);
964    }
965}
olai_http/lib.rs

olai_http/
lib.rs
