Merge branch 'main' into feature-parity-with-ytproxy

.github/workflows/build.yml

@@ -18,12 +18,12 @@ jobs:
       - uses: Swatinem/rust-cache@v2
       - uses: rui314/setup-mold@v1
       - name: Set up NASM
-        uses: ilammy/setup-nasm@v1.4.0
+        uses: ilammy/setup-nasm@v1.5.1
       - name: Build
         run: RUSTFLAGS='-C target-feature=+crt-static' cargo build --release --target x86_64-unknown-linux-gnu
       - run: mv target/x86_64-unknown-linux-gnu/release/piped-proxy piped-proxy
       - name: Upload artifact
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: piped-proxy
           path: piped-proxy

Cargo.toml

@@ -7,25 +7,25 @@ version = "0.1.0"
 
 [dependencies]
 # Web Requests & Async Runtime
-tokio = { version = "1.34.0", features = ["full"] }
-actix-web = "4.4.0"
-reqwest = { version = "0.11.22", features = ["stream", "brotli", "gzip", "socks"], default-features = false }
+tokio = { version = "1.36.0", features = ["full"] }
+actix-web = "4.5.1"
+reqwest = { version = "0.11.25", features = ["stream", "brotli", "gzip", "socks"], default-features = false }
 qstring = "0.7.2"
 
 # Alternate Allocator 
 mimalloc = { version = "0.1.39", optional = true }
 
 # Transcoding Images to WebP/AVIF to save bandwidth
-image = { version = "0.24.7", features = ["jpeg", "jpeg_rayon", "webp"], default-features = false, optional = true }
-libwebp-sys = { version = "0.9.4", optional = true }
-ravif = { version = "0.11.3", optional = true }
+image = { version = "0.24.9", features = ["jpeg", "jpeg_rayon", "webp"], default-features = false, optional = true }
+libwebp-sys = { version = "0.9.5", optional = true }
+ravif = { version = "0.11.4", optional = true }
 rgb = { version = "0.8.37", optional = true }
 
-once_cell = "1.18.0"
-regex = "1.10.2"
+once_cell = "1.19.0"
+regex = "1.10.3"
 blake3 = { version = "1.5.0", optional = true }
 bytes = "1.5.0"
-futures-util = "0.3.29"
+futures-util = "0.3.30"
 
 [features]
 default = ["webp", "mimalloc", "reqwest-rustls", "qhash"]

src/main.rs

@@ -1,24 +1,23 @@
-use actix_web::http::Method;
+mod ump_stream;
+mod utils;
+
+use actix_web::http::{Method, StatusCode};
 use actix_web::{web, App, HttpRequest, HttpResponse, HttpResponseBuilder, HttpServer};
 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;
 use std::io::ErrorKind;
-use std::pin::Pin;
-use std::task::{Context, Poll};
 use std::{env, io};
 
 #[cfg(not(any(feature = "reqwest-native-tls", feature = "reqwest-rustls")))]
 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::TryStreamExt;
 #[cfg(any(feature = "webp", feature = "avif", feature = "qhash"))]
 use tokio::task::spawn_blocking;
+use ump_stream::UmpTransformStream;
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]
@@ -47,7 +46,6 @@ async fn main() -> std::io::Result<()> {
     .await
 }
 
-static PREFIX_PATH: Lazy<String> = Lazy::new(|| String::from(env::var("PREFIX_PATH").unwrap_or_else(|_| "".to_string())));
 static RE_DOMAIN: Lazy<Regex> =
     Lazy::new(|| Regex::new(r"^(?:[a-z\d.-]*\.)?([a-z\d-]*\.[a-z\d-]*)$").unwrap());
 static RE_MANIFEST: Lazy<Regex> = Lazy::new(|| Regex::new("(?m)URI=\"([^\"]+)\"").unwrap());
@@ -98,7 +96,6 @@ const ALLOWED_DOMAINS: [&str; 8] = [
     "ajay.app",
 ];
 
-
 fn add_headers(response: &mut HttpResponseBuilder) {
     response
         .append_header(("Access-Control-Allow-Origin", "*"))
@@ -114,7 +111,7 @@ fn is_header_allowed(header: &str) -> bool {
 
     !matches!(
         header,
-            | "host"
+        "host"
             | "authorization"
             | "origin"
             | "referer"
@@ -127,6 +124,8 @@ fn is_header_allowed(header: &str) -> bool {
             | "report-to"
             | "strict-transport-security"
             | "user-agent"
+            | "range"
+            | "transfer-encoding"
     )
 }
 
@@ -191,7 +190,20 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
                     hasher.update(&value);
                 }
 
-                hasher.update(&path);
+                let range_marker = b"/range/";
+
+                // Find the slice before "/range/"
+                if let Some(position) = path
+                    .windows(range_marker.len())
+                    .position(|window| window == range_marker)
+                {
+                    // Update the hasher with the part of the path before "/range/"
+                    // We add +1 to include the "/" in the hash
+                    // This is done for DASH streams for the manifests provided by YouTube
+                    hasher.update(&path[..(position + 1)]);
+                } else {
+                    hasher.update(&path);
+                }
 
                 hasher.update(secret.as_bytes());
 
@@ -243,14 +255,40 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
 
     let mime_type = query.get("mime").map(|s| s.to_string());
 
-    if is_ump && !query.has("range") {
+    let clen = query
+        .get("clen")
+        .map(|s| s.to_string().parse::<u64>().unwrap());
+
+    if video_playback && !query.has("range") {
         if let Some(range) = req.headers().get("range") {
             let range = range.to_str().unwrap();
             let range = range.replace("bytes=", "");
-            query.add_pair(("range", range));
+            let range = range.split('-').collect::<Vec<_>>();
+            let start = range[0].parse::<u64>().unwrap();
+            let end = match range[1].parse::<u64>() {
+                Ok(end) => end,
+                Err(_) => {
+                    if let Some(clen) = clen {
+                        clen - 1
+                    } else {
+                        0
+                    }
+                }
+            };
+            if end != 0 {
+                let range = format!("{}-{}", start, end);
+                query.add_pair(("range", range));
+            }
+        } else {
+            if let Some(clen) = clen {
+                let range = format!("0-{}", clen - 1);
+                query.add_pair(("range", range));
+            }
         }
     }
 
+    let range = query.get("range").map(|s| s.to_string());
+
     let qs = {
         let collected = query
             .into_pairs()
@@ -399,11 +437,13 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
                         if let Some(captures) = captures {
                             let url = captures.get(1).unwrap().as_str();
                             if url.starts_with("https://") {
-                                return line
-                                    .replace(url, localize_url(url, host.as_str()).as_str());
+                                return line.replace(
+                                    url,
+                                    utils::localize_url(url, host.as_str()).as_str(),
+                                );
                             }
                         }
-                        localize_url(line, host.as_str())
+                        utils::localize_url(line, host.as_str())
                     })
                     .collect::<Vec<String>>()
                     .join("\n");
@@ -416,8 +456,9 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
                 let captures = RE_DASH_MANIFEST.captures_iter(&resp_str);
                 for capture in captures {
                     let url = capture.get(1).unwrap().as_str();
-                    let new_url = localize_url(url, host.as_str());
-                    new_resp = new_resp.replace(url, new_url.as_str());
+                    let new_url = utils::localize_url(url, host.as_str());
+                    let new_url = utils::escape_xml(new_url.as_str());
+                    new_resp = new_resp.replace(url, new_url.as_ref());
                 }
                 return Ok(response.body(new_resp));
             }
@@ -425,228 +466,49 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
     }
 
     if let Some(content_length) = resp.headers().get("content-length") {
-        response.append_header(("content-length", content_length));
+        response.no_chunking(content_length.to_str().unwrap().parse::<u64>().unwrap());
     }
 
-    let resp = resp.bytes_stream();
-
-    if is_ump {
+    if is_ump && resp.status().is_success() {
         if let Some(mime_type) = mime_type {
             response.content_type(mime_type);
         }
+        if req.headers().contains_key("range") {
+            // check if it's not the whole stream
+            if let Some(ref range) = range {
+                if let Some(clen) = clen {
+                    if range != &format!("0-{}", clen - 1) {
+                        response.status(StatusCode::PARTIAL_CONTENT);
+                    }
+                }
+            }
+        }
+        let resp = resp.bytes_stream();
+        let resp = resp.map_err(|e| io::Error::new(ErrorKind::Other, e));
         let transformed_stream = UmpTransformStream::new(resp);
+        // print errors
+        let transformed_stream = transformed_stream.map_err(|e| {
+            eprintln!("UMP Transforming Error: {}", e);
+            e
+        });
+
+        // calculate content length from clen and range
+        if let Some(clen) = clen {
+            let length = if let Some(ref range) = range {
+                let range = range.replace("bytes=", "");
+                let range = range.split('-').collect::<Vec<_>>();
+                let start = range[0].parse::<u64>().unwrap();
+                let end = range[1].parse::<u64>().unwrap_or(clen - 1);
+                end - start + 1
+            } else {
+                clen
+            };
+            response.no_chunking(length);
+        }
+
         return Ok(response.streaming(transformed_stream));
     }
 
     // Stream response
-    Ok(response.streaming(resp))
-}
-
-fn read_buf(buf: &[u8], pos: &mut usize) -> u8 {
-    let byte = buf[*pos];
-    *pos += 1;
-    byte
-}
-
-fn read_variable_integer(buf: &[u8], offset: usize) -> io::Result<(i32, usize)> {
-    let mut pos = offset;
-    let prefix = read_buf(buf, &mut pos);
-    let mut size = 0;
-    for shift in 1..=5 {
-        if prefix & (128 >> (shift - 1)) == 0 {
-            size = shift;
-            break;
-        }
-    }
-    if !(1..=5).contains(&size) {
-        return Err(io::Error::new(
-            ErrorKind::InvalidData,
-            format!("Invalid integer size {} at position {}", size, offset),
-        ));
-    }
-
-    match size {
-        1 => Ok((prefix as i32, size)),
-        2 => {
-            let value = ((read_buf(buf, &mut pos) as i32) << 6) | (prefix as i32 & 0b111111);
-            Ok((value, size))
-        }
-        3 => {
-            let value =
-                (((read_buf(buf, &mut pos) as i32) | ((read_buf(buf, &mut pos) as i32) << 8)) << 5)
-                    | (prefix as i32 & 0b11111);
-            Ok((value, size))
-        }
-        4 => {
-            let value = (((read_buf(buf, &mut pos) as i32)
-                | ((read_buf(buf, &mut pos) as i32) << 8)
-                | ((read_buf(buf, &mut pos) as i32) << 16))
-                << 4)
-                | (prefix as i32 & 0b1111);
-            Ok((value, size))
-        }
-        _ => {
-            let value = (read_buf(buf, &mut pos) as i32)
-                | ((read_buf(buf, &mut pos) as i32) << 8)
-                | ((read_buf(buf, &mut pos) as i32) << 16)
-                | ((read_buf(buf, &mut pos) as i32) << 24);
-            Ok((value, size))
-        }
-    }
-}
-
-struct UmpTransformStream<S>
-where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
-{
-    inner: S,
-    buffer: BytesMut,
-    found_stream: bool,
-    remaining: usize,
-}
-
-impl<S> UmpTransformStream<S>
-where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
-{
-    pub fn new(stream: S) -> Self {
-        UmpTransformStream {
-            inner: stream,
-            buffer: BytesMut::new(),
-            found_stream: false,
-            remaining: 0,
-        }
-    }
-}
-
-impl<S> Stream for UmpTransformStream<S>
-where
-    S: Stream<Item = Result<Bytes, ReqwestError>> + Unpin,
-{
-    type Item = Result<Bytes, ReqwestError>;
-
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.get_mut();
-
-        while let Poll::Ready(item) = Pin::new(&mut this.inner).poll_next(cx) {
-            match item {
-                Some(Ok(bytes)) => {
-                    if this.found_stream {
-                        if this.remaining > 0 {
-                            let len = std::cmp::min(this.remaining, bytes.len());
-                            this.remaining -= len;
-                            if this.remaining == 0 {
-                                this.buffer.clear();
-                                this.buffer.extend_from_slice(&bytes[len..]);
-                                this.found_stream = false;
-                            }
-                            return Poll::Ready(Some(Ok(bytes.slice(0..len))));
-                        } else {
-                            this.found_stream = false;
-                            this.buffer.clear();
-                            this.buffer.extend_from_slice(&bytes);
-                        };
-                    } else {
-                        this.buffer.extend_from_slice(&bytes);
-                    }
-                }
-                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
-                None => {
-                    return Poll::Ready(None);
-                }
-            }
-        }
-
-        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();
-            if segment_type != 21 {
-                // Not the stream
-                if this.buffer.len() > s1 + s2 + segment_length as usize {
-                    let _ = this.buffer.split_to(s1 + s2 + segment_length as usize);
-                }
-            } else {
-                this.remaining = segment_length as usize - 1;
-
-                let _ = this.buffer.split_to(s1 + s2 + 1);
-
-                if this.buffer.len() > segment_length as usize {
-                    let len = std::cmp::min(this.remaining, this.buffer.len());
-                    this.remaining -= len;
-
-                    return Poll::Ready(Some(Ok(this.buffer.split_to(len).into())));
-                } else {
-                    this.remaining -= this.buffer.len();
-                    this.found_stream = true;
-
-                    return Poll::Ready(Some(Ok(this.buffer.to_vec().into())));
-                }
-            }
-        }
-
-        Poll::Pending
-    }
-}
-
-fn finalize_url(path: &str, query: BTreeMap<String, String>) -> String {
-    #[cfg(feature = "qhash")]
-    {
-        use std::collections::BTreeSet;
-
-        let qhash = {
-            let secret = env::var("HASH_SECRET");
-            if let Ok(secret) = secret {
-                let set = query
-                    .iter()
-                    .filter(|(key, _)| !matches!(key.as_str(), "qhash" | "range" | "rewrite"))
-                    .map(|(key, value)| (key.as_bytes().to_owned(), value.as_bytes().to_owned()))
-                    .collect::<BTreeSet<_>>();
-
-                let mut hasher = blake3::Hasher::new();
-
-                for (key, value) in set {
-                    hasher.update(&key);
-                    hasher.update(&value);
-                }
-
-                hasher.update(path.as_bytes());
-
-                hasher.update(secret.as_bytes());
-
-                let hash = hasher.finalize().to_hex();
-
-                Some(hash[..8].to_owned())
-            } else {
-                None
-            }
-        };
-
-        if qhash.is_some() {
-            let mut query = QString::new(query.into_iter().collect::<Vec<_>>());
-            query.add_pair(("qhash", qhash.unwrap()));
-            return format!("{}{}?{}", PREFIX_PATH.as_str(), path, query);
-        }
-    }
-
-    let query = QString::new(query.into_iter().collect::<Vec<_>>());
-    format!("{}{}?{}", PREFIX_PATH.as_str(), path, query)
-}
-
-fn localize_url(url: &str, host: &str) -> String {
-    if url.starts_with("https://") {
-        let url = Url::parse(url).unwrap();
-        let host = url.host().unwrap().to_string();
-
-        let mut query = url.query_pairs().into_owned().collect::<BTreeMap<_, _>>();
-        query.insert("host".to_string(), host.clone());
-
-        return finalize_url(url.path(), query);
-    } else if url.ends_with(".m3u8") || url.ends_with(".ts") {
-        let mut query = BTreeMap::new();
-        query.insert("host".to_string(), host.to_string());
-
-        return finalize_url(url, query);
-    }
-
-    url.to_string()
-}
+    Ok(response.streaming(resp.bytes_stream()))
+}

src/ump_stream.rs

@@ -0,0 +1,154 @@
+use crate::utils;
+use bytes::{Bytes, BytesMut};
+use futures_util::Stream;
+use std::io;
+use std::io::ErrorKind;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+fn read_variable_integer(buf: &[u8], offset: usize) -> io::Result<(i32, usize)> {
+    let mut pos = offset;
+    let prefix = utils::read_buf(buf, &mut pos);
+    let mut size = 0;
+    for shift in 1..=5 {
+        if prefix & (128 >> (shift - 1)) == 0 {
+            size = shift;
+            break;
+        }
+    }
+    if !(1..=5).contains(&size) {
+        return Err(io::Error::new(
+            ErrorKind::InvalidData,
+            format!("Invalid integer size {} at position {}", size, offset),
+        ));
+    }
+
+    match size {
+        1 => Ok((prefix as i32, size)),
+        2 => {
+            let value = ((utils::read_buf(buf, &mut pos) as i32) << 6) | (prefix as i32 & 0b111111);
+            Ok((value, size))
+        }
+        3 => {
+            let value = (((utils::read_buf(buf, &mut pos) as i32)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 8))
+                << 5)
+                | (prefix as i32 & 0b11111);
+            Ok((value, size))
+        }
+        4 => {
+            let value = (((utils::read_buf(buf, &mut pos) as i32)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 8)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 16))
+                << 4)
+                | (prefix as i32 & 0b1111);
+            Ok((value, size))
+        }
+        _ => {
+            let value = (utils::read_buf(buf, &mut pos) as i32)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 8)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 16)
+                | ((utils::read_buf(buf, &mut pos) as i32) << 24);
+            Ok((value, size))
+        }
+    }
+}
+
+pub struct UmpTransformStream<S>
+where
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
+{
+    inner: S,
+    buffer: BytesMut,
+    found_stream: bool,
+    remaining: usize,
+}
+
+impl<S> UmpTransformStream<S>
+where
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
+{
+    pub fn new(stream: S) -> Self {
+        UmpTransformStream {
+            inner: stream,
+            buffer: BytesMut::new(),
+            found_stream: false,
+            remaining: 0,
+        }
+    }
+}
+
+impl<S> Stream for UmpTransformStream<S>
+where
+    S: Stream<Item = Result<Bytes, io::Error>> + Unpin,
+{
+    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();
+
+        while let Poll::Ready(item) = Pin::new(&mut this.inner).poll_next(cx) {
+            match item {
+                Some(Ok(bytes)) => {
+                    if this.found_stream {
+                        if this.remaining > 0 {
+                            let len = std::cmp::min(this.remaining, bytes.len());
+                            this.remaining -= len;
+                            if this.remaining == 0 {
+                                this.buffer.clear();
+                                this.buffer.extend_from_slice(&bytes[len..]);
+                                this.found_stream = false;
+                            }
+                            return Poll::Ready(Some(Ok(bytes.slice(0..len))));
+                        } else {
+                            this.found_stream = false;
+                            this.buffer.clear();
+                            this.buffer.extend_from_slice(&bytes);
+                        };
+                    } else {
+                        this.buffer.extend_from_slice(&bytes);
+                    }
+                }
+                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+                None => {
+                    return Poll::Ready(None);
+                }
+            }
+        }
+
+        if !this.found_stream && !this.buffer.is_empty() {
+            let (segment_type, s1) = match read_variable_integer(&this.buffer, 0) {
+                Ok(result) => result,
+                Err(_) => return Poll::Pending,
+            };
+            let (segment_length, s2) = match read_variable_integer(&this.buffer, s1) {
+                Ok(result) => result,
+                Err(_) => return Poll::Pending,
+            };
+            if segment_type != 21 {
+                // Not the stream
+                if this.buffer.len() > s1 + s2 + segment_length as usize {
+                    let _ = this.buffer.split_to(s1 + s2 + segment_length as usize);
+                }
+            } else {
+                this.remaining = segment_length as usize - 1;
+
+                let _ = this.buffer.split_to(s1 + s2 + 1);
+
+                if this.buffer.len() > segment_length as usize {
+                    let len = std::cmp::min(this.remaining, this.buffer.len());
+                    this.remaining -= len;
+
+                    return Poll::Ready(Some(Ok(this.buffer.split_to(len).into())));
+                } else {
+                    this.remaining -= this.buffer.len();
+                    this.found_stream = true;
+
+                    return Poll::Ready(Some(Ok(this.buffer.to_vec().into())));
+                }
+            }
+        }
+
+        Poll::Pending
+    }
+}

src/utils.rs

@@ -0,0 +1,100 @@
+use once_cell::sync::Lazy;
+use qstring::QString;
+use reqwest::Url;
+use std::borrow::Cow;
+use std::collections::BTreeMap;
+use std::env;
+
+static PREFIX_PATH: Lazy<String> = Lazy::new(|| String::from(env::var("PREFIX_PATH").unwrap_or_else(|_| "".to_string())));
+
+pub fn read_buf(buf: &[u8], pos: &mut usize) -> u8 {
+    let byte = buf[*pos];
+    *pos += 1;
+    byte
+}
+
+fn finalize_url(path: &str, query: BTreeMap<String, String>) -> String {
+    #[cfg(feature = "qhash")]
+    {
+        use std::collections::BTreeSet;
+        use std::env;
+
+        let qhash = {
+            let secret = env::var("HASH_SECRET");
+            if let Ok(secret) = secret {
+                let set = query
+                    .iter()
+                    .filter(|(key, _)| !matches!(key.as_str(), "qhash" | "range" | "rewrite"))
+                    .map(|(key, value)| (key.as_bytes().to_owned(), value.as_bytes().to_owned()))
+                    .collect::<BTreeSet<_>>();
+
+                let mut hasher = blake3::Hasher::new();
+
+                for (key, value) in set {
+                    hasher.update(&key);
+                    hasher.update(&value);
+                }
+
+                hasher.update(path.as_bytes());
+
+                hasher.update(secret.as_bytes());
+
+                let hash = hasher.finalize().to_hex();
+
+                Some(hash[..8].to_owned())
+            } else {
+                None
+            }
+        };
+
+        if qhash.is_some() {
+            let mut query = QString::new(query.into_iter().collect::<Vec<_>>());
+            query.add_pair(("qhash", qhash.unwrap()));
+            return format!("{}{}?{}", PREFIX_PATH.as_str(), path, query);
+        }
+    }
+
+    let query = QString::new(query.into_iter().collect::<Vec<_>>());
+    format!("{}{}?{}", PREFIX_PATH.as_str(), path, query)
+}
+
+pub fn localize_url(url: &str, host: &str) -> String {
+    if url.starts_with("https://") {
+        let url = Url::parse(url).unwrap();
+        let host = url.host().unwrap().to_string();
+
+        let mut query = url.query_pairs().into_owned().collect::<BTreeMap<_, _>>();
+
+        query.insert("host".to_string(), host.clone());
+
+        return finalize_url(url.path(), query);
+    } else if url.ends_with(".m3u8") || url.ends_with(".ts") {
+        let mut query = BTreeMap::new();
+        query.insert("host".to_string(), host.to_string());
+
+        return finalize_url(url, query);
+    }
+
+    url.to_string()
+}
+
+pub fn escape_xml(raw: &str) -> Cow<'_, str> {
+    if !raw.contains(&['<', '>', '&', '\'', '"'][..]) {
+        // If there are no characters to escape, return the original string.
+        Cow::Borrowed(raw)
+    } else {
+        // If there are characters to escape, build a new string with the replacements.
+        let mut escaped = String::with_capacity(raw.len());
+        for c in raw.chars() {
+            match c {
+                '<' => escaped.push_str("&lt;"),
+                '>' => escaped.push_str("&gt;"),
+                '&' => escaped.push_str("&amp;"),
+                '\'' => escaped.push_str("&apos;"),
+                '"' => escaped.push_str("&quot;"),
+                _ => escaped.push(c),
+            }
+        }
+        Cow::Owned(escaped)
+    }
+}