use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::PathBuf;
use std::sync::Arc;
use tauri::http::{header, Method, Request, Response, StatusCode};
use tauri::plugin::Builder as PluginBuilder;
use tauri::{Manager, Runtime};
use super::allow::Allow;
use super::read::{read_head, sniff, sniffed_content_type, HEAD_BYTES};
use crate::Shared;
pub const SCHEME: &str = "drim-file";
const RANGE_CHUNK: u64 = 4 * 1024 * 1024;
const FULL_LIMIT: u64 = 64 * 1024 * 1024;
pub fn register<R: Runtime>(builder: PluginBuilder<R>) -> PluginBuilder<R> {
builder.register_asynchronous_uri_scheme_protocol(SCHEME, |ctx, request, responder| {
let app = ctx.app_handle().clone();
let label = ctx.webview_label().to_string();
tauri::async_runtime::spawn_blocking(move || {
let response = match app.try_state::<Arc<Shared>>() {
Some(shared) if shared.hub.get(&label).is_some() => serve(&shared.files.allow, &request),
Some(_) => {
log::debug!("drim-file:窗口 {label} 没有 hello,拒绝");
status(StatusCode::FORBIDDEN)
}
None => status(StatusCode::SERVICE_UNAVAILABLE),
};
responder.respond(response);
});
})
}
#[allow(unused_variables)]
pub fn enable_cors<R: Runtime>(webview: &tauri::Webview<R>) {
#[cfg(target_os = "linux")]
{
let _ = webview.with_webview(|platform| {
use webkit2gtk::{SecurityManagerExt, WebContextExt, WebViewExt};
match platform.inner().context().and_then(|c| c.security_manager()) {
Some(manager) => {
if !manager.uri_scheme_is_cors_enabled("drim-file") {
manager.register_uri_scheme_as_cors_enabled("drim-file");
log::info!("drim-file 已登记为允许跨域");
}
}
None => log::warn!("取不到 WebKitGTK 的 SecurityManager,drim-file 不能跨域读取"),
}
});
}
}
fn base(code: StatusCode) -> tauri::http::response::Builder {
Response::builder()
.status(code)
.header(header::ACCESS_CONTROL_ALLOW_ORIGIN, "*")
.header(header::ACCESS_CONTROL_EXPOSE_HEADERS, "Content-Range, Content-Length, Accept-Ranges")
.header(header::X_CONTENT_TYPE_OPTIONS, "nosniff")
.header(header::CACHE_CONTROL, "no-cache")
}
fn status(code: StatusCode) -> Response<Vec<u8>> {
base(code).body(Vec::new()).unwrap()
}
pub fn path_of(uri: &tauri::http::Uri) -> Option<PathBuf> {
let raw = uri.path().strip_prefix('/').unwrap_or(uri.path());
if raw.is_empty() {
return None;
}
let bytes: Vec<u8> = percent_encoding::percent_decode_str(raw).collect();
#[cfg(unix)]
{
use std::os::unix::ffi::OsStrExt;
let path = PathBuf::from(std::ffi::OsStr::from_bytes(&bytes));
Some(if path.is_absolute() { path } else { PathBuf::from("/").join(path) })
}
#[cfg(not(unix))]
{
Some(PathBuf::from(String::from_utf8(bytes).ok()?))
}
}
fn parse_range(value: &str, size: u64) -> Option<Option<(u64, u64)>> {
let spec = value.trim().strip_prefix("bytes=")?;
let (a, b) = spec.split(',').next()?.trim().split_once('-')?;
if a.is_empty() {
let n: u64 = b.trim().parse().ok()?;
if n == 0 || size == 0 {
return Some(None);
}
return Some(Some((size.saturating_sub(n), size - 1)));
}
let start: u64 = a.trim().parse().ok()?;
if start >= size {
return Some(None);
}
let end = if b.trim().is_empty() { size - 1 } else { b.trim().parse::<u64>().ok()?.min(size - 1) };
if end < start {
return None;
}
Some(Some((start, end)))
}
pub fn serve(allow: &Allow, request: &Request<Vec<u8>>) -> Response<Vec<u8>> {
let method = request.method();
if method == Method::OPTIONS {
return base(StatusCode::NO_CONTENT)
.header(header::ACCESS_CONTROL_ALLOW_METHODS, "GET, HEAD, OPTIONS")
.header(header::ACCESS_CONTROL_ALLOW_HEADERS, "Range")
.body(Vec::new())
.unwrap();
}
if method != Method::GET && method != Method::HEAD {
return status(StatusCode::METHOD_NOT_ALLOWED);
}
let Some(path) = path_of(request.uri()) else { return status(StatusCode::BAD_REQUEST) };
let real = match allow.check_read(&path) {
Ok(real) => real,
Err(err) if err.reason.as_deref() == Some("file_not_found") => return status(StatusCode::NOT_FOUND),
Err(_) => return status(StatusCode::FORBIDDEN),
};
let Ok(mut file) = File::open(&real) else { return status(StatusCode::FORBIDDEN) };
let size = match file.metadata() {
Ok(meta) if meta.is_file() => meta.len(),
_ => return status(StatusCode::NOT_FOUND),
};
let head = read_head(&mut file, HEAD_BYTES).unwrap_or_default();
let content_type = sniffed_content_type(sniff(&head));
let range = request.headers().get(header::RANGE).and_then(|v| v.to_str().ok()).and_then(|v| parse_range(v, size));
let (code, start, end) = match range {
Some(None) => {
return base(StatusCode::RANGE_NOT_SATISFIABLE).header(header::CONTENT_RANGE, format!("bytes */{size}")).body(Vec::new()).unwrap();
}
Some(Some((start, end))) => (StatusCode::PARTIAL_CONTENT, start, end.min(start + RANGE_CHUNK - 1)),
None if size > FULL_LIMIT => (StatusCode::PARTIAL_CONTENT, 0, RANGE_CHUNK - 1),
None => (StatusCode::OK, 0, size.saturating_sub(1)),
};
let length = if size == 0 { 0 } else { end - start + 1 };
let mut builder = base(code).header(header::CONTENT_TYPE, content_type).header(header::ACCEPT_RANGES, "bytes").header(header::CONTENT_LENGTH, length);
if code == StatusCode::PARTIAL_CONTENT {
builder = builder.header(header::CONTENT_RANGE, format!("bytes {start}-{end}/{size}"));
}
if method == Method::HEAD {
return builder.body(Vec::new()).unwrap();
}
let mut body = Vec::with_capacity(length as usize);
let read = file.seek(SeekFrom::Start(start)).and_then(|_| file.by_ref().take(length).read_to_end(&mut body));
if read.is_err() || body.len() as u64 != length {
return status(StatusCode::INTERNAL_SERVER_ERROR);
}
builder.body(body).unwrap()
}
#[cfg(test)]
mod tests {
use super::*;
fn get(uri: &str, range: Option<&str>) -> Request<Vec<u8>> {
let mut builder = Request::builder().uri(uri);
if let Some(r) = range {
builder = builder.header(header::RANGE, r);
}
builder.body(Vec::new()).unwrap()
}
fn url_of(path: &std::path::Path) -> String {
let encoded = percent_encoding::utf8_percent_encode(path.to_str().unwrap(), percent_encoding::NON_ALPHANUMERIC).to_string();
format!("drim-file://localhost/{encoded}")
}
#[test]
fn serves_allowed_files_with_ranges() {
let dir = tempfile::tempdir().unwrap();
let temp = dir.path().join("temp");
std::fs::create_dir_all(&temp).unwrap();
let allow = Allow::new(dir.path().join("data"), temp.clone(), None, vec![], vec![]);
let video = temp.join("a b.mp4");
let mut bytes = b"\0\0\0\x18ftypisom".to_vec();
bytes.resize(1000, 7);
std::fs::write(&video, &bytes).unwrap();
let url = url_of(&video);
let full = serve(&allow, &get(&url, None));
assert_eq!(full.status(), StatusCode::OK);
assert_eq!(full.headers()[header::CONTENT_TYPE], "video/mp4");
assert_eq!(full.headers()[header::ACCESS_CONTROL_ALLOW_ORIGIN], "*");
assert_eq!(full.body().len(), 1000);
let part = serve(&allow, &get(&url, Some("bytes=100-199")));
assert_eq!(part.status(), StatusCode::PARTIAL_CONTENT);
assert_eq!(part.headers()[header::CONTENT_RANGE], "bytes 100-199/1000");
assert_eq!(part.body().as_slice(), &bytes[100..200]);
let open_ended = serve(&allow, &get(&url, Some("bytes=990-")));
assert_eq!(open_ended.body().len(), 10);
let suffix = serve(&allow, &get(&url, Some("bytes=-4")));
assert_eq!(suffix.headers()[header::CONTENT_RANGE], "bytes 996-999/1000");
assert_eq!(serve(&allow, &get(&url, Some("bytes=5000-"))).status(), StatusCode::RANGE_NOT_SATISFIABLE);
#[cfg(unix)]
{
let plain = format!("drim-file://localhost{}", video.to_str().unwrap().replace(' ', "%20"));
assert_eq!(serve(&allow, &get(&plain, Some("bytes=0-0"))).status(), StatusCode::PARTIAL_CONTENT);
}
let outside = dir.path().join("secret.txt");
std::fs::write(&outside, b"<html><script>").unwrap();
assert_eq!(serve(&allow, &get(&url_of(&outside), None)).status(), StatusCode::FORBIDDEN);
assert_eq!(serve(&allow, &get(&url_of(&temp.join("none.mp4")), None)).status(), StatusCode::NOT_FOUND);
let html = temp.join("page.html");
std::fs::write(&html, b"<html><script>alert(1)</script>").unwrap();
assert_eq!(serve(&allow, &get(&url_of(&html), None)).headers()[header::CONTENT_TYPE], "application/octet-stream");
}
}