salvo_proxy/
reqwest_client.rs

1use futures_util::TryStreamExt;
2use hyper::upgrade::OnUpgrade;
3use reqwest::Client as InnerClient;
4use salvo_core::http::{ResBody, StatusCode};
5use salvo_core::rt::tokio::TokioIo;
6use salvo_core::Error;
7use tokio::io::copy_bidirectional;
8
9use crate::{Client, HyperRequest, BoxedError, Proxy, Upstreams, HyperResponse};
10
11/// A [`Client`] implementation based on [`reqwest::Client`].
12/// 
13/// This client provides proxy capabilities using the Reqwest HTTP client.
14/// It supports all features of Reqwest including automatic redirect handling,
15/// connection pooling, and other HTTP client features.
16#[derive(Default, Clone, Debug)]
17pub struct ReqwestClient {
18    inner: InnerClient,
19}
20
21impl<U> Proxy<U, ReqwestClient>
22where
23    U: Upstreams,
24    U::Error: Into<BoxedError>,
25{
26    /// Create a new `Proxy` using the default Reqwest client.
27    /// 
28    /// This is a convenient way to create a proxy with standard configuration.
29    pub fn use_reqwest_client(upstreams: U) -> Self {
30        Proxy::new(upstreams, ReqwestClient::default())
31    }
32}
33
34impl ReqwestClient {
35    /// Create a new `ReqwestClient` with the given [`reqwest::Client`].
36    pub fn new(inner: InnerClient) -> Self {
37        Self { inner }
38    }
39}
40
41impl Client for ReqwestClient {
42    type Error = salvo_core::Error;
43
44    async fn execute(
45        &self,
46        proxied_request: HyperRequest,
47        request_upgraded: Option<OnUpgrade>,
48    ) -> Result<HyperResponse, Self::Error> {
49        let request_upgrade_type = crate::get_upgrade_type(proxied_request.headers()).map(|s| s.to_owned());
50
51        let proxied_request =
52            proxied_request.map(|s| reqwest::Body::wrap_stream(s.map_ok(|s| s.into_data().unwrap_or_default())));
53        let response = self
54            .inner
55            .execute(proxied_request.try_into().map_err(Error::other)?)
56            .await
57            .map_err(Error::other)?;
58
59        let res_headers = response.headers().clone();
60        let hyper_response = hyper::Response::builder()
61            .status(response.status())
62            .version(response.version());
63
64        let mut hyper_response = if response.status() == StatusCode::SWITCHING_PROTOCOLS {
65            let response_upgrade_type = crate::get_upgrade_type(response.headers());
66
67            if request_upgrade_type == response_upgrade_type.map(|s| s.to_lowercase()) {
68                let mut response_upgraded = response
69                    .upgrade()
70                    .await
71                    .map_err(|e| Error::other(format!("response does not have an upgrade extension. {}", e)))?;
72                if let Some(request_upgraded) = request_upgraded {
73                    tokio::spawn(async move {
74                        match request_upgraded.await {
75                            Ok(request_upgraded) => {
76                                let mut request_upgraded = TokioIo::new(request_upgraded);
77                                if let Err(e) = copy_bidirectional(&mut response_upgraded, &mut request_upgraded).await
78                                {
79                                    tracing::error!(error = ?e, "coping between upgraded connections failed");
80                                }
81                            }
82                            Err(e) => {
83                                tracing::error!(error = ?e, "upgrade request failed");
84                            }
85                        }
86                    });
87                } else {
88                    return Err(Error::other("request does not have an upgrade extension"));
89                }
90            } else {
91                return Err(Error::other("upgrade type mismatch"));
92            }
93            hyper_response.body(ResBody::None).map_err(Error::other)?
94        } else {
95            hyper_response
96                .body(ResBody::stream(response.bytes_stream()))
97                .map_err(Error::other)?
98        };
99        *hyper_response.headers_mut() = res_headers;
100        Ok(hyper_response)
101    }
102}
103
104// Unit tests for Proxy
105#[cfg(test)]
106mod tests {
107    use salvo_core::prelude::*;
108    use salvo_core::test::*;
109
110    use super::*;
111    use crate::{Upstreams, Proxy};
112
113    #[tokio::test]
114    async fn test_upstreams_elect() {
115        let upstreams = vec!["https://www.example.com", "https://www.example2.com"];
116        let proxy = Proxy::new(upstreams.clone(), ReqwestClient::default());
117        let elected_upstream = proxy.upstreams().elect().await.unwrap();
118        assert!(upstreams.contains(&elected_upstream));
119    }
120
121    #[tokio::test]
122    async fn test_reqwest_client() {
123        let router = Router::new().push(
124            Router::with_path("rust/{**rest}").goal(Proxy::new(vec!["https://www.rust-lang.org"], ReqwestClient::default())),
125        );
126
127        let content = TestClient::get("http://127.0.0.1:5801/rust/tools/install")
128            .send(router)
129            .await
130            .take_string()
131            .await
132            .unwrap();
133        assert!(content.contains("Install Rust"));
134    }
135
136    #[test]
137    fn test_others() {
138        let mut handler = Proxy::new(["https://www.bing.com"], ReqwestClient::default());
139        assert_eq!(handler.upstreams().len(), 1);
140        assert_eq!(handler.upstreams_mut().len(), 1);
141    }
142}
salvo_proxy/reqwest_client.rs

salvo_proxy/
reqwest_client.rs
