mirror of
				https://github.com/TeamPiped/piped-proxy.git
				synced 2024-08-14 23:50:45 +00:00 
			
		
		
		
	Merge branch 'main' into feature-parity-with-ytproxy
This commit is contained in:
		
						commit
						84c790a154
					
				
					 6 changed files with 808 additions and 580 deletions
				
			
		
							
								
								
									
										4
									
								
								.github/workflows/build.yml
									
										
									
									
										vendored
									
									
								
							
							
						
						
									
										4
									
								
								.github/workflows/build.yml
									
										
									
									
										vendored
									
									
								
							| 
						 | 
					@ -18,12 +18,12 @@ jobs:
 | 
				
			||||||
      - uses: Swatinem/rust-cache@v2
 | 
					      - uses: Swatinem/rust-cache@v2
 | 
				
			||||||
      - uses: rui314/setup-mold@v1
 | 
					      - uses: rui314/setup-mold@v1
 | 
				
			||||||
      - name: Set up NASM
 | 
					      - name: Set up NASM
 | 
				
			||||||
        uses: ilammy/setup-nasm@v1.4.0
 | 
					        uses: ilammy/setup-nasm@v1.5.1
 | 
				
			||||||
      - name: Build
 | 
					      - name: Build
 | 
				
			||||||
        run: RUSTFLAGS='-C target-feature=+crt-static' cargo build --release --target x86_64-unknown-linux-gnu
 | 
					        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
 | 
					      - run: mv target/x86_64-unknown-linux-gnu/release/piped-proxy piped-proxy
 | 
				
			||||||
      - name: Upload artifact
 | 
					      - name: Upload artifact
 | 
				
			||||||
        uses: actions/upload-artifact@v3
 | 
					        uses: actions/upload-artifact@v4
 | 
				
			||||||
        with:
 | 
					        with:
 | 
				
			||||||
          name: piped-proxy
 | 
					          name: piped-proxy
 | 
				
			||||||
          path: piped-proxy
 | 
					          path: piped-proxy
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
							
								
								
									
										784
									
								
								Cargo.lock
									
										
									
										generated
									
									
									
								
							
							
						
						
									
										784
									
								
								Cargo.lock
									
										
									
										generated
									
									
									
								
							
										
											
												File diff suppressed because it is too large
												Load diff
											
										
									
								
							
							
								
								
									
										18
									
								
								Cargo.toml
									
										
									
									
									
								
							
							
						
						
									
										18
									
								
								Cargo.toml
									
										
									
									
									
								
							| 
						 | 
					@ -7,25 +7,25 @@ version = "0.1.0"
 | 
				
			||||||
 | 
					
 | 
				
			||||||
[dependencies]
 | 
					[dependencies]
 | 
				
			||||||
# Web Requests & Async Runtime
 | 
					# Web Requests & Async Runtime
 | 
				
			||||||
tokio = { version = "1.34.0", features = ["full"] }
 | 
					tokio = { version = "1.36.0", features = ["full"] }
 | 
				
			||||||
actix-web = "4.4.0"
 | 
					actix-web = "4.5.1"
 | 
				
			||||||
reqwest = { version = "0.11.22", features = ["stream", "brotli", "gzip", "socks"], default-features = false }
 | 
					reqwest = { version = "0.11.25", features = ["stream", "brotli", "gzip", "socks"], default-features = false }
 | 
				
			||||||
qstring = "0.7.2"
 | 
					qstring = "0.7.2"
 | 
				
			||||||
 | 
					
 | 
				
			||||||
# Alternate Allocator 
 | 
					# Alternate Allocator 
 | 
				
			||||||
mimalloc = { version = "0.1.39", optional = true }
 | 
					mimalloc = { version = "0.1.39", optional = true }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
# Transcoding Images to WebP/AVIF to save bandwidth
 | 
					# Transcoding Images to WebP/AVIF to save bandwidth
 | 
				
			||||||
image = { version = "0.24.7", features = ["jpeg", "jpeg_rayon", "webp"], default-features = false, optional = true }
 | 
					image = { version = "0.24.9", features = ["jpeg", "jpeg_rayon", "webp"], default-features = false, optional = true }
 | 
				
			||||||
libwebp-sys = { version = "0.9.4", optional = true }
 | 
					libwebp-sys = { version = "0.9.5", optional = true }
 | 
				
			||||||
ravif = { version = "0.11.3", optional = true }
 | 
					ravif = { version = "0.11.4", optional = true }
 | 
				
			||||||
rgb = { version = "0.8.37", optional = true }
 | 
					rgb = { version = "0.8.37", optional = true }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
once_cell = "1.18.0"
 | 
					once_cell = "1.19.0"
 | 
				
			||||||
regex = "1.10.2"
 | 
					regex = "1.10.3"
 | 
				
			||||||
blake3 = { version = "1.5.0", optional = true }
 | 
					blake3 = { version = "1.5.0", optional = true }
 | 
				
			||||||
bytes = "1.5.0"
 | 
					bytes = "1.5.0"
 | 
				
			||||||
futures-util = "0.3.29"
 | 
					futures-util = "0.3.30"
 | 
				
			||||||
 | 
					
 | 
				
			||||||
[features]
 | 
					[features]
 | 
				
			||||||
default = ["webp", "mimalloc", "reqwest-rustls", "qhash"]
 | 
					default = ["webp", "mimalloc", "reqwest-rustls", "qhash"]
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
							
								
								
									
										322
									
								
								src/main.rs
									
										
									
									
									
								
							
							
						
						
									
										322
									
								
								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 actix_web::{web, App, HttpRequest, HttpResponse, HttpResponseBuilder, HttpServer};
 | 
				
			||||||
use once_cell::sync::Lazy;
 | 
					use once_cell::sync::Lazy;
 | 
				
			||||||
use qstring::QString;
 | 
					use qstring::QString;
 | 
				
			||||||
use regex::Regex;
 | 
					use regex::Regex;
 | 
				
			||||||
use reqwest::Error as ReqwestError;
 | 
					 | 
				
			||||||
use reqwest::{Body, Client, Request, Url};
 | 
					use reqwest::{Body, Client, Request, Url};
 | 
				
			||||||
use std::collections::BTreeMap;
 | 
					 | 
				
			||||||
use std::error::Error;
 | 
					use std::error::Error;
 | 
				
			||||||
use std::io::ErrorKind;
 | 
					use std::io::ErrorKind;
 | 
				
			||||||
use std::pin::Pin;
 | 
					 | 
				
			||||||
use std::task::{Context, Poll};
 | 
					 | 
				
			||||||
use std::{env, io};
 | 
					use std::{env, io};
 | 
				
			||||||
 | 
					
 | 
				
			||||||
#[cfg(not(any(feature = "reqwest-native-tls", feature = "reqwest-rustls")))]
 | 
					#[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");
 | 
					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::TryStreamExt;
 | 
				
			||||||
use futures_util::Stream;
 | 
					 | 
				
			||||||
#[cfg(any(feature = "webp", feature = "avif", feature = "qhash"))]
 | 
					#[cfg(any(feature = "webp", feature = "avif", feature = "qhash"))]
 | 
				
			||||||
use tokio::task::spawn_blocking;
 | 
					use tokio::task::spawn_blocking;
 | 
				
			||||||
 | 
					use ump_stream::UmpTransformStream;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
#[cfg(feature = "mimalloc")]
 | 
					#[cfg(feature = "mimalloc")]
 | 
				
			||||||
#[global_allocator]
 | 
					#[global_allocator]
 | 
				
			||||||
| 
						 | 
					@ -47,7 +46,6 @@ async fn main() -> std::io::Result<()> {
 | 
				
			||||||
    .await
 | 
					    .await
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
static PREFIX_PATH: Lazy<String> = Lazy::new(|| String::from(env::var("PREFIX_PATH").unwrap_or_else(|_| "".to_string())));
 | 
					 | 
				
			||||||
static RE_DOMAIN: Lazy<Regex> =
 | 
					static RE_DOMAIN: Lazy<Regex> =
 | 
				
			||||||
    Lazy::new(|| Regex::new(r"^(?:[a-z\d.-]*\.)?([a-z\d-]*\.[a-z\d-]*)$").unwrap());
 | 
					    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());
 | 
					static RE_MANIFEST: Lazy<Regex> = Lazy::new(|| Regex::new("(?m)URI=\"([^\"]+)\"").unwrap());
 | 
				
			||||||
| 
						 | 
					@ -98,7 +96,6 @@ const ALLOWED_DOMAINS: [&str; 8] = [
 | 
				
			||||||
    "ajay.app",
 | 
					    "ajay.app",
 | 
				
			||||||
];
 | 
					];
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					 | 
				
			||||||
fn add_headers(response: &mut HttpResponseBuilder) {
 | 
					fn add_headers(response: &mut HttpResponseBuilder) {
 | 
				
			||||||
    response
 | 
					    response
 | 
				
			||||||
        .append_header(("Access-Control-Allow-Origin", "*"))
 | 
					        .append_header(("Access-Control-Allow-Origin", "*"))
 | 
				
			||||||
| 
						 | 
					@ -114,7 +111,7 @@ fn is_header_allowed(header: &str) -> bool {
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    !matches!(
 | 
					    !matches!(
 | 
				
			||||||
        header,
 | 
					        header,
 | 
				
			||||||
            | "host"
 | 
					        "host"
 | 
				
			||||||
            | "authorization"
 | 
					            | "authorization"
 | 
				
			||||||
            | "origin"
 | 
					            | "origin"
 | 
				
			||||||
            | "referer"
 | 
					            | "referer"
 | 
				
			||||||
| 
						 | 
					@ -127,6 +124,8 @@ fn is_header_allowed(header: &str) -> bool {
 | 
				
			||||||
            | "report-to"
 | 
					            | "report-to"
 | 
				
			||||||
            | "strict-transport-security"
 | 
					            | "strict-transport-security"
 | 
				
			||||||
            | "user-agent"
 | 
					            | "user-agent"
 | 
				
			||||||
 | 
					            | "range"
 | 
				
			||||||
 | 
					            | "transfer-encoding"
 | 
				
			||||||
    )
 | 
					    )
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					@ -191,7 +190,20 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
 | 
				
			||||||
                    hasher.update(&value);
 | 
					                    hasher.update(&value);
 | 
				
			||||||
                }
 | 
					                }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					                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(&path);
 | 
				
			||||||
 | 
					                }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
                hasher.update(secret.as_bytes());
 | 
					                hasher.update(secret.as_bytes());
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					@ -243,13 +255,39 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    let mime_type = query.get("mime").map(|s| s.to_string());
 | 
					    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") {
 | 
					        if let Some(range) = req.headers().get("range") {
 | 
				
			||||||
            let range = range.to_str().unwrap();
 | 
					            let range = range.to_str().unwrap();
 | 
				
			||||||
            let range = range.replace("bytes=", "");
 | 
					            let range = range.replace("bytes=", "");
 | 
				
			||||||
 | 
					            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));
 | 
					                query.add_pair(("range", range));
 | 
				
			||||||
            }
 | 
					            }
 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    let range = query.get("range").map(|s| s.to_string());
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    let qs = {
 | 
					    let qs = {
 | 
				
			||||||
        let collected = query
 | 
					        let collected = query
 | 
				
			||||||
| 
						 | 
					@ -399,11 +437,13 @@ async fn index(req: HttpRequest) -> Result<HttpResponse, Box<dyn Error>> {
 | 
				
			||||||
                        if let Some(captures) = captures {
 | 
					                        if let Some(captures) = captures {
 | 
				
			||||||
                            let url = captures.get(1).unwrap().as_str();
 | 
					                            let url = captures.get(1).unwrap().as_str();
 | 
				
			||||||
                            if url.starts_with("https://") {
 | 
					                            if url.starts_with("https://") {
 | 
				
			||||||
                                return line
 | 
					                                return line.replace(
 | 
				
			||||||
                                    .replace(url, localize_url(url, host.as_str()).as_str());
 | 
					                                    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>>()
 | 
					                    .collect::<Vec<String>>()
 | 
				
			||||||
                    .join("\n");
 | 
					                    .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);
 | 
					                let captures = RE_DASH_MANIFEST.captures_iter(&resp_str);
 | 
				
			||||||
                for capture in captures {
 | 
					                for capture in captures {
 | 
				
			||||||
                    let url = capture.get(1).unwrap().as_str();
 | 
					                    let url = capture.get(1).unwrap().as_str();
 | 
				
			||||||
                    let new_url = localize_url(url, host.as_str());
 | 
					                    let new_url = utils::localize_url(url, host.as_str());
 | 
				
			||||||
                    new_resp = new_resp.replace(url, new_url.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));
 | 
					                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") {
 | 
					    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 && resp.status().is_success() {
 | 
				
			||||||
 | 
					 | 
				
			||||||
    if is_ump {
 | 
					 | 
				
			||||||
        if let Some(mime_type) = mime_type {
 | 
					        if let Some(mime_type) = mime_type {
 | 
				
			||||||
            response.content_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);
 | 
					        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));
 | 
					        return Ok(response.streaming(transformed_stream));
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    // Stream response
 | 
					    // Stream response
 | 
				
			||||||
    Ok(response.streaming(resp))
 | 
					    Ok(response.streaming(resp.bytes_stream()))
 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
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()
 | 
					 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
							
								
								
									
										154
									
								
								src/ump_stream.rs
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										154
									
								
								src/ump_stream.rs
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
					@ -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
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
							
								
								
									
										100
									
								
								src/utils.rs
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										100
									
								
								src/utils.rs
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
					@ -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("<"),
 | 
				
			||||||
 | 
					                '>' => escaped.push_str(">"),
 | 
				
			||||||
 | 
					                '&' => escaped.push_str("&"),
 | 
				
			||||||
 | 
					                '\'' => escaped.push_str("'"),
 | 
				
			||||||
 | 
					                '"' => escaped.push_str("""),
 | 
				
			||||||
 | 
					                _ => escaped.push(c),
 | 
				
			||||||
 | 
					            }
 | 
				
			||||||
 | 
					        }
 | 
				
			||||||
 | 
					        Cow::Owned(escaped)
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue