aria2-core 0.2.2

High-performance download engine core: multi-protocol segmented downloads, rate limiting, config management, session persistence, and BitTorrent seeding
Documentation
use crate::error::Result;
use async_trait::async_trait;
use std::any::Any;
use std::path::Path;

#[async_trait]
pub trait DiskAdaptor: Send + Sync {
    async fn open(&mut self, path: &Path) -> Result<()>;
    async fn write(&mut self, offset: u64, data: &[u8]) -> Result<()>;
    async fn read(&mut self, offset: u64, length: u64) -> Result<Vec<u8>>;
    async fn close(&mut self) -> Result<()>;
    async fn truncate(&mut self, length: u64) -> Result<()>;
    async fn flush(&mut self) -> Result<()>;
    async fn size(&self) -> Result<u64>;
    fn as_any(&self) -> &dyn Any;

    #[cfg(unix)]
    fn unix_raw_fd(&self) -> Option<std::os::unix::io::RawFd>;

    /// Returns the raw OS file handle on Windows, or `None` if no file is open.
    /// The handle is borrowed (not owned); callers must not close it.
    #[cfg(windows)]
    fn windows_raw_handle(&self) -> Option<std::os::windows::io::RawHandle>;
}

pub struct DirectDiskAdaptor {
    file: Option<tokio::fs::File>,
    path: std::path::PathBuf,
}

impl DirectDiskAdaptor {
    pub fn new() -> Self {
        DirectDiskAdaptor {
            file: None,
            path: std::path::PathBuf::new(),
        }
    }
}

impl Default for DirectDiskAdaptor {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl DiskAdaptor for DirectDiskAdaptor {
    async fn open(&mut self, path: &Path) -> Result<()> {
        self.path = path.to_path_buf();
        let mut open_opts = tokio::fs::OpenOptions::new();

        if path.exists() {
            open_opts.write(true).read(true);
        } else {
            open_opts.write(true).create(true).read(true);
        }

        self.file = Some(
            open_opts
                .open(path)
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?,
        );

        Ok(())
    }

    async fn write(&mut self, offset: u64, data: &[u8]) -> Result<()> {
        if let Some(ref mut file) = self.file {
            use tokio::io::{AsyncSeekExt, AsyncWriteExt};
            file.seek(std::io::SeekFrom::Start(offset))
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;
            file.write_all(data)
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;
        }
        Ok(())
    }

    async fn read(&mut self, offset: u64, length: u64) -> Result<Vec<u8>> {
        if let Some(ref mut file) = self.file {
            use tokio::io::{AsyncReadExt, AsyncSeekExt};
            file.seek(std::io::SeekFrom::Start(offset))
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;

            let mut buffer = vec![0u8; length as usize];
            let bytes_read = file.read_exact(&mut buffer).await;

            match bytes_read {
                Ok(_) => Ok(buffer),
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::UnexpectedEof {
                        Ok(buffer)
                    } else {
                        Err(crate::error::Aria2Error::Io(e.to_string()))
                    }
                }
            }
        } else {
            Err(crate::error::Aria2Error::DownloadFailed(
                "文件未打开".to_string(),
            ))
        }
    }

    async fn close(&mut self) -> Result<()> {
        self.file = None;
        Ok(())
    }

    async fn truncate(&mut self, length: u64) -> Result<()> {
        if let Some(ref mut file) = self.file {
            file.set_len(length)
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;
        }
        Ok(())
    }

    async fn flush(&mut self) -> Result<()> {
        if let Some(ref mut file) = self.file {
            use tokio::io::AsyncWriteExt;
            file.flush()
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;
        }
        Ok(())
    }

    async fn size(&self) -> Result<u64> {
        if let Some(ref file) = self.file {
            let metadata = file
                .metadata()
                .await
                .map_err(|e| crate::error::Aria2Error::Io(e.to_string()))?;
            Ok(metadata.len())
        } else {
            Err(crate::error::Aria2Error::DownloadFailed(
                "文件未打开".to_string(),
            ))
        }
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    #[cfg(unix)]
    fn unix_raw_fd(&self) -> Option<std::os::unix::io::RawFd> {
        use std::os::fd::AsRawFd;
        self.file.as_ref().map(|f| f.as_raw_fd())
    }

    #[cfg(windows)]
    fn windows_raw_handle(&self) -> Option<std::os::windows::io::RawHandle> {
        use std::os::windows::io::AsRawHandle;
        self.file.as_ref().map(|f| f.as_raw_handle())
    }
}