From 7586ae314b2012f1f7729b2c6d27927faa3ae327 Mon Sep 17 00:00:00 2001
From: Kavin <20838718+FireMasterK@users.noreply.github.com>
Date: Mon, 11 Dec 2023 22:51:12 +0000
Subject: [PATCH] Refactor UmpTransformStream to use io::Error instead of
 ReqwestError

- Replaced all occurrences of `ReqwestError` with `io::Error` in the code.
- Updated the type signature and implementation of `UmpTransformStream` to use `io::Error`.
- Handled errors when calling `read_variable_integer` by returning an error result.
---
 src/main.rs | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 51b0698..c1150ea 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,6 @@ use actix_web::{web, App, HttpRequest, HttpResponse, HttpResponseBuilder, HttpSe
 use once_cell::sync::Lazy;
 use qstring::QString;
 use regex::Regex;
-use reqwest::Error as ReqwestError;
 use reqwest::{Body, Client, Request, Url};
 use std::collections::BTreeMap;
 use std::error::Error;
@@ -16,7 +15,7 @@ use std::{env, io};
 compile_error!("feature \"reqwest-native-tls\" or \"reqwest-rustls\" must be set for proxy to have TLS support");
 
 use bytes::{Bytes, BytesMut};
-use futures_util::Stream;
+use futures_util::{Stream, TryStreamExt};
 #[cfg(any(feature = "webp", feature = "avif", feature = "qhash"))]
 use tokio::task::spawn_blocking;
 
@@ -431,6 +430,7 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
         if req.headers().contains_key("range") {
             response.status(StatusCode::PARTIAL_CONTENT);
         }
+        let resp = resp.map_err(|e| io::Error::new(ErrorKind::Other, e));
         let transformed_stream = UmpTransformStream::new(resp);
         return Ok(response.streaming(transformed_stream));
     }
@@ -494,7 +494,7 @@ fn read_variable_integer(buf: &[u8], offset: usize) -> io::Result<(i32, usize)>
 
 struct UmpTransformStream<S>
 where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
 {
     inner: S,
     buffer: BytesMut,
@@ -504,7 +504,7 @@ where
 
 impl<S> UmpTransformStream<S>
 where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
 {
     pub fn new(stream: S) -> Self {
         UmpTransformStream {
@@ -518,9 +518,9 @@ where
 
 impl<S> Stream for UmpTransformStream<S>
 where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
 {
-    type Item = Result<Bytes, ReqwestError>;
+    type Item = Result<Bytes, io::Error>;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.get_mut();
@@ -555,8 +555,14 @@ where
         }
 
         if !this.found_stream && !this.buffer.is_empty() {
-            let (segment_type, s1) = read_variable_integer(&this.buffer, 0).unwrap();
-            let (segment_length, s2) = read_variable_integer(&this.buffer, s1).unwrap();
+            let (segment_type, s1) = match read_variable_integer(&this.buffer, 0) {
+                Ok(result) => result,
+                Err(e) => return Poll::Ready(Some(Err(e))),
+            };
+            let (segment_length, s2) = match read_variable_integer(&this.buffer, s1) {
+                Ok(result) => result,
+                Err(e) => return Poll::Ready(Some(Err(e))),
+            };
             if segment_type != 21 {
                 // Not the stream
                 if this.buffer.len() > s1 + s2 + segment_length as usize {