refactor: Update reverse_proxy.yaml with new upstream locations

arloor · Sep 20, 2024 · 7bf7594 · 7bf7594
1 parent 00c61df
commit 7bf7594
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 19 deletions.
diff --git a/rust_http_proxy/src/proxy.rs b/rust_http_proxy/src/proxy.rs
@@ -407,7 +407,7 @@ fn mod_http1_proxy_req(req: &mut Request<Incoming>) -> io::Result<()> {
         info!("change host header: {:?} -> {:?}", origin, host_header);
     }
     // change absoulte uri to relative uri
-    origin_form(req.uri_mut());
+    origin_form(req.uri_mut())?;
     Ok(())
 }
 
@@ -746,19 +746,20 @@ pub(crate) fn check_auth(
     (username, authed)
 }
 
-fn origin_form(uri: &mut Uri) {
+fn origin_form(uri: &mut Uri) -> io::Result<()> {
     let path = match uri.path_and_query() {
         Some(path) if path.as_str() != "/" => {
             let mut parts = ::http::uri::Parts::default();
             parts.path_and_query = Some(path.clone());
-            Uri::from_parts(parts).expect("path is valid uri")
+            Uri::from_parts(parts).map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?
         }
         _none_or_just_slash => {
             debug_assert!(Uri::default() == "/");
             Uri::default()
         }
     };
-    *uri = path
+    *uri = path;
+    Ok(())
 }
 
 // Create a TCP connection to host:port, build a tunnel between the connection and

diff --git a/rust_http_proxy/src/web_func.rs b/rust_http_proxy/src/web_func.rs
@@ -221,7 +221,13 @@ async fn serve_metrics(
             .header(http::header::CONTENT_TYPE, "text/plain; charset=utf-8")
             .header(http::header::SERVER, SERVER_NAME);
         if can_gzip {
-            let compressed_data = compress_string(&buffer);
+            let compressed_data = match compress_string(&buffer) {
+                Ok(compressed_data) => compressed_data,
+                Err(e) => {
+                    warn!("compress metrics error: {}", e);
+                    return Ok(build_500_resp());
+                }
+            };
             builder
                 .header(http::header::CONTENT_ENCODING, GZIP)
                 .body(full_body(compressed_data))
@@ -561,7 +567,13 @@ async fn _speed(
         .header(http::header::SERVER, SERVER_NAME)
         .header(http::header::CONTENT_TYPE, "text/html; charset=utf-8");
     if can_gzip {
-        let compressed_data = compress_string(&body);
+        let compressed_data = match compress_string(&body) {
+            Ok(compressed_data) => compressed_data,
+            Err(e) => {
+                warn!("compress body error: {}", e);
+                return Ok(build_500_resp());
+            }
+        };
         builder
             .header(http::header::CONTENT_ENCODING, GZIP)
             .body(full_body(compressed_data))
@@ -580,24 +592,20 @@ use flate2::write::GzEncoder;
 use flate2::Compression;
 use std::io::prelude::*;
 
-fn compress_string(input: &str) -> Vec<u8> {
+fn compress_string(input: &str) -> io::Result<Vec<u8>> {
     let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
-    encoder
-        .write_all(input.as_bytes())
-        .expect("Failed to write data");
-    encoder.finish().expect("Failed to finish compression")
+    encoder.write_all(input.as_bytes())?;
+    encoder.finish()
 }
 
 #[allow(unused)]
 use flate2::read::GzDecoder;
 #[allow(unused)]
-fn decompress_string(input: &[u8]) -> String {
+fn decompress_string(input: &[u8]) -> io::Result<String> {
     let mut decoder = GzDecoder::new(input);
     let mut decompressed_data = String::new();
-    decoder
-        .read_to_string(&mut decompressed_data)
-        .expect("Failed to read data");
-    decompressed_data
+    decoder.read_to_string(&mut decompressed_data)?;
+    Ok(decompressed_data)
 }
 
 #[cfg(test)]
@@ -643,11 +651,12 @@ mod tests {
     }
 
     #[test]
-    fn test_gzip_compress_string() {
+    fn test_gzip_compress_string() -> io::Result<()> {
         let original_string = "Hello, Rust! This is a test string for Gzip compression.";
-        let compressed_data = compress_string(original_string);
-        let decompressed_string = decompress_string(&compressed_data);
+        let compressed_data = compress_string(original_string)?;
+        let decompressed_string = decompress_string(&compressed_data)?;
 
         assert_eq!(original_string, decompressed_string);
+        Ok(())
     }
 }