Rust 返回文件流接口

如果是直接返回整个文件,可以直接使用 Axum提供的 ReadSteam 构造一个文件流返回就可以。

async fn file_stream_handler() -> Response {
    // 打开文件
    let file = File::open("test.log").await.unwrap();

    // 将文件转换为异步流
    let stream = ReaderStream::new(file);

    // 创建 StreamBody
    let body = StreamBody::new(stream);
        
    // 设置响应头
    let mut headers= header::HeaderMap::new();
    headers.insert(header::CONTENT_TYPE, header::HeaderValue::from_static("application/octet-stream"));
    
    // 返回响应
    body.into_response()
}

#[tokio::main]
async fn main() {
    // 构建路由
    let app = Router::new().route("/stream-file", get(file_stream_handler));

    // 绑定地址
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    println!("Listening on {}", addr);

    // 启动服务器
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

但是有可能是需要文件中的一部分,或者要更细致控制流,就不能这么简单的使用了,可以自己实现一个流,通过seek 位移来实现读取文件某一块

async fn file_stream_handler(Path(path): Path<String>) -> Response {
    // 构造自定义的文件流
    let file_stream = FileStream::new(path.as_str(), 0, 0, 1024).await;

    // 设置响应头
    let mut headers = header::HeaderMap::new();
    headers.insert(
        header::CONTENT_TYPE,
        header::HeaderValue::from_static("application/octet-stream"),
    );

    (headers, StreamBody::new(file_stream)).into_response()
}

struct FileStream {
    reader: File,
    end: u64,
    current: u64,
    buffer: Vec<u8>,
}

impl FileStream {
    async fn new(file_path:&str, start: u64, mut end: u64, buffer_size: usize) -> Self {
        let mut file = File::open(file_path).await.unwrap();
        file.seek(std::io::SeekFrom::Start(start)).await.unwrap();

        if end == 0 {
            let metadata = file.metadata().await.unwrap();
            end = metadata.len() as u64;
        }
    
        FileStream {
            reader: file,
            end,
            current: 0,
            buffer: vec![0; buffer_size],
        }
    }
}

impl Stream for FileStream {
    type Item = Result<Bytes, std::io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();

        // 确保当前指针在读取范围内
        if this.current >= this.end {
            return Poll::Ready(None); // 已读取完指定范围
        }

        // 计算要读取的字节数
        let bytes_to_read = std::cmp::min(this.buffer.len() as u64, this.end - this.current);
        let bytes_to_read = bytes_to_read as usize;

        // 使用 poll_read 方法从文件中读取数据
        let mut read_buf = ReadBuf::new(&mut this.buffer[..bytes_to_read]);
        match Pin::new(&mut this.reader).poll_read(cx, &mut read_buf) {
            Poll::Ready(Ok(())) => {
                let n = read_buf.filled().len();
                if n == 0 {
                    return Poll::Ready(None); // 文件结束
                }
                this.current += n as u64; // 更新当前读取的字节数
                Poll::Ready(Some(Ok(Bytes::copy_from_slice(&this.buffer[..n]))))
            }
            Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))), // 读取错误
            Poll::Pending => Poll::Pending, // 读取尚未完成
        }
    }
}

async_stream 提供了一个快速实现流的宏,可以更快捷的实现

async fn file_stream_handler(Path(path): Path<String>) -> Response {
    // 构造自定义的文件流
    //let file_stream = FileStream::new(path.as_str(), 0, 0, 1024).await;
    let file_stream = try_stream(path.as_str(), 0, 0, 1024).await;

    // 设置响应头
    let mut headers = header::HeaderMap::new();
    headers.insert(
        header::CONTENT_TYPE,
        header::HeaderValue::from_static("application/octet-stream"),
    );

    (headers, StreamBody::new(file_stream)).into_response()
}


async fn try_stream(
    file_path: &str,
    start: u64,
    mut end: u64,
    buffer_size: usize,
) -> impl Stream<Item = Result<Vec<u8>, std::io::Error>> {
    // 构造文件
    let mut file = File::open(file_path).await.unwrap();
    file.seek(std::io::SeekFrom::Start(start)).await.unwrap();
    
    // 没规定截取点就读完
    if end == 0 {
        let metadata = file.metadata().await.unwrap();
        end = metadata.len() as u64;
    }
   
    let mut buffer = vec![0; buffer_size];
    // 使用 async-stream 的 try_stream! 宏快捷实现流返回
    let stream = try_stream! {
        let mut total_read = 0;

            while total_read < end {
                // 比较还剩下多少没有读完,取小的部分
                let bytes_to_read = std::cmp::min(buffer_size as u64, end - total_read);
                
                let n = file.read(&mut buffer[..bytes_to_read as usize]).await.map_err(|e| {
                    std::io::Error::new(std::io::ErrorKind::Other, e)
                })?;

                if n == 0 {
                    break; // EOF
                }
                
                total_read += n as u64;
                yield buffer[..n].to_vec(); // 返回读取的字节

        }
    };
    stream
}