Manager: 修复/public接口不支持返回206 partial content的问题

BalloonUpdate · Jan 31, 2025 · 73828be · 73828be
1 parent 99b579f
commit 73828be
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 19 deletions.
diff --git a/manager/src/builtin_server/mod.rs b/manager/src/builtin_server/mod.rs
@@ -161,7 +161,7 @@ async fn _send_data(stream: &mut TcpStream, data: &[u8]) -> std::io::Result<()>
     Ok(())
 }
 
-async fn receive_data<'a>(stream: &'a mut TcpStream) -> std::io::Result<PartialAsyncRead<'a, TcpStream>> {
+async fn receive_data<'a>(stream: &'a mut TcpStream) -> std::io::Result<PartialAsyncRead<&'a mut TcpStream>> {
     let len = timeout(stream.read_u64_le()).await?;
 
     Ok(PartialAsyncRead::new(stream, len))

diff --git a/manager/src/common/tar_reader.rs b/manager/src/common/tar_reader.rs
@@ -31,7 +31,7 @@ impl TarReader {
     }
 
     /// 读取更新包中的一个文件数据，需要提供文件的`offset`和`len`以便定位
-    pub fn open_file(&mut self, offset: u64, len: u64) -> PartialRead<std::fs::File> {
+    pub fn open_file(&mut self, offset: u64, len: u64) -> PartialRead<&mut std::fs::File> {
         self.open.seek(SeekFrom::Start(offset)).unwrap();
 
         PartialRead::new(&mut self.open, len)

diff --git a/manager/src/web/api/public/mod.rs b/manager/src/web/api/public/mod.rs
@@ -1,12 +1,16 @@
 use axum::body::Body;
 use axum::extract::Path;
 use axum::extract::State;
+use axum::http::HeaderMap;
 use axum::response::Response;
+use reqwest::StatusCode;
 use shared::utility::filename_ext::GetFileNamePart;
+use shared::utility::partial_read::PartialAsyncRead;
+use tokio::io::AsyncSeekExt;
 
 use crate::web::webstate::WebState;
 
-pub async fn api_public(State(state): State<WebState>, Path(path): Path<String>) -> Response {
+pub async fn api_public(State(state): State<WebState>, headers: HeaderMap, Path(path): Path<String>) -> Response {
     println!("+public: {}", path);
 
     let path = state.app_path.public_dir.join(path);
@@ -15,20 +19,58 @@ pub async fn api_public(State(state): State<WebState>, Path(path): Path<String>)
         return Response::builder().status(404).body(Body::empty()).unwrap();
     }
 
+    let range = headers.get("range")
+        // 拿出range的值
+        .map(|e| e.to_str().unwrap())
+        // 检查是否以bytes开头
+        .filter(|e| e.starts_with("bytes="))
+        // 提取出bytes=后面的部分
+        .map(|e| e["bytes=".len()..].split("-"))
+        // 提取出开始字节和结束字节
+        .and_then(|mut e| Some((e.next()?, e.next()?)))
+        // 解析开始字节和结束字节
+        .and_then(|e| Some((u64::from_str_radix(e.0, 10).ok()?, u64::from_str_radix(e.1, 10).ok()? + 1)))
+        // 开始和结束不能都等于0
+        .filter(|e| e != &(0, 0))
+        // 转换成range
+        .map(|e| e.0..e.1);
+
+    // 检查range参数
+    if let Some(range) = &range {
+        if range.end < range.start {
+            return Response::builder().status(403).body(Body::from("incorrect range")).unwrap();
+        }
+    }
+
     let metadata = tokio::fs::metadata(&path).await.unwrap();
 
-    let file = tokio::fs::File::options()
+    let mut file = tokio::fs::File::options()
         .read(true)
         .open(&path)
         .await
         .unwrap();
 
-    let file = tokio_util::io::ReaderStream::new(file);
+    if let Some(range) = &range {
+        file.seek(std::io::SeekFrom::Start(range.start)).await.unwrap();
+    }
+
+    let len = match &range {
+        Some(range) => range.end - range.start,
+        None => metadata.len(),
+    };
+
+    let file = tokio_util::io::ReaderStream::new(PartialAsyncRead::new(file, len));
+
+    let mut builder = Response::builder();
+
+    builder = builder.header(axum::http::header::CONTENT_TYPE, "application/octet-stream");
+    builder = builder.header(axum::http::header::CONTENT_DISPOSITION, format!("attachment; filename=\"{}\"", path.filename()));
+    builder = builder.header(axum::http::header::CONTENT_LENGTH, format!("{}", len));
+
+    if let Some(range) = &range {
+        builder = builder.header(axum::http::header::CONTENT_RANGE, format!("{}-{}/{}", range.start, range.end - 1, metadata.len()));
+        builder = builder.status(StatusCode::PARTIAL_CONTENT);
+    }
 
-    Response::builder()
-        .header(axum::http::header::CONTENT_TYPE, "application/octet-stream")
-        .header(axum::http::header::CONTENT_DISPOSITION, format!("attachment; filename=\"{}\"", path.filename()))
-        .header(axum::http::header::CONTENT_LENGTH, format!("{}", metadata.len()))
-        .body(Body::from_stream(file))
-        .unwrap()
+    builder.body(Body::from_stream(file)).unwrap()
 }
diff --git a/shared/src/utility/partial_read.rs b/shared/src/utility/partial_read.rs
@@ -4,15 +4,15 @@ use tokio::io::AsyncRead;
 use tokio::pin;
 
 /// 代表一个限制读取数量的Read
-pub struct PartialRead<'a, R: Read>(&'a mut R, u64);
+pub struct PartialRead<R: Read>(R, u64);
 
-impl<'a, R: Read> PartialRead<'a, R> {
-    pub fn new(read: &'a mut R, count: u64) -> Self {
+impl<R: Read> PartialRead<R> {
+    pub fn new(read: R, count: u64) -> Self {
         Self(read, count)
     }
 }
 
-impl<R: Read> Read for PartialRead<'_, R> {
+impl<R: Read> Read for PartialRead<R> {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
         if self.1 == 0 {
             return Ok(0);
@@ -24,10 +24,10 @@ impl<R: Read> Read for PartialRead<'_, R> {
 }
 
 /// 代表一个限制读取数量的Read
-pub struct PartialAsyncRead<'a, R: AsyncRead>(&'a mut R, u64);
+pub struct PartialAsyncRead<R: AsyncRead>(R, u64);
 
-impl<'a, R: AsyncRead> PartialAsyncRead<'a, R> {
-    pub fn new(read: &'a mut R, count: u64) -> Self {
+impl<R: AsyncRead> PartialAsyncRead<R> {
+    pub fn new(read: R, count: u64) -> Self {
         Self(read, count)
     }
 
@@ -36,7 +36,7 @@ impl<'a, R: AsyncRead> PartialAsyncRead<'a, R> {
     }
 }
 
-impl<R: AsyncRead + Unpin> AsyncRead for PartialAsyncRead<'_, R> {
+impl<R: AsyncRead + Unpin> AsyncRead for PartialAsyncRead<R> {
     fn poll_read(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,