rlist_drivers/

onedrive.rs

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use rlist_driver_macro::{StaticCombinableFile, StaticDownloadLinkFile, VfsMeta};
use rlist_vfs::combinable_dir::CombinableDir;
use rlist_vfs::driver::{CloudDriver, GetVfs};
use rlist_vfs::static_combinable::StaticCombinableFile;
use rlist_vfs::static_combinable::StaticDownloadLinkFile;
use rlist_vfs::VfsBasicMeta;
use serde::Deserialize;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::SystemTime;
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

#[derive(Debug, Deserialize, Clone)]
pub struct OnedriveConfig {
    /// The refresh token for the onedrive account.
    /// *For further information, please refer to the official documentation of Microsoft OAuth 2.0 authorization flow.*
    pub refresh_token: String,

    /// The client id for the application.
    /// You can get it from the Azure portal with the client secret.
    pub client_id: String,

    /// The client secret for the application.
    /// You can get it from the Azure portal with the client id.
    pub client_secret: String,
}

pub struct OnedriveState {
    pub config: OnedriveConfig,
    pub access_token: Arc<Mutex<String>>,
    pub expires_at: Arc<Mutex<i64>>,
    pub my_drive_id: String,
}

pub struct OnedriveDriver {
    pub state: OnedriveState,
}

#[async_trait]
impl GetVfs for OnedriveDriver {
    async fn get_vfs(&self) -> Result<CombinableDir<StaticCombinableFile>, String> {
        Self::reload_vfs(&self.state).await
    }
}

#[async_trait]
impl CloudDriver<OnedriveConfig, OnedriveState> for OnedriveDriver {
    async fn new(state: OnedriveState) -> Self {
        Self { state }
    }

    async fn load_config(config: OnedriveConfig) -> OnedriveState {
        let access_token = fetch_access_token(&config).await.unwrap();
        let AccessTokenResponse {
            access_token,
            expires_in,
            ..
        } = access_token;
        let my_drive_id = get_my_od_id(&access_token).await.unwrap();
        OnedriveState {
            config,
            access_token: Arc::new(Mutex::new(access_token)),
            expires_at: Arc::new(Mutex::new(expires_in)),
            my_drive_id,
        }
    }

    async fn reload_vfs(
        state: &OnedriveState,
    ) -> Result<CombinableDir<StaticCombinableFile>, String> {
        info!(
            "Onedrive Driver: Reloading VFS for drive {}.",
            state.my_drive_id.clone()
        );
        // judge whether the access token is expired
        let mut expires_at_lock = state.expires_at.lock().await;
        let expired = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs()
            > *expires_at_lock as u64;
        let mut access_token_lock = state.access_token.lock().await;

        // if expired, refresh the access token
        if expired {
            debug!(
                "Onedrive Driver: Access token of drive {} is expired, refreshing.",
                state.my_drive_id.clone()
            );
            let new_access_token = fetch_access_token(&state.config).await?;
            let AccessTokenResponse {
                access_token,
                expires_in,
                ..
            } = new_access_token;
            *access_token_lock = access_token;
            *expires_at_lock = expires_in;
            debug!(
                "Onedrive Driver: Access token of drive {} is refreshed.",
                state.my_drive_id.clone()
            );
        } // end if expired

        // request and build the tree
        debug!(
            "Onedrive Driver: Building the tree for drive {}.",
            state.my_drive_id.clone()
        );
        let tree = build_tree(
            access_token_lock.clone(),
            state.my_drive_id.clone(),
            "root".to_owned(),
            "root".to_owned(),
            0,
            SystemTime::now(),
        )
        .await;
        debug!(
            "Onedrive Driver: Tree for drive {} is built.",
            state.my_drive_id.clone()
        );

        let (folder, error_count) = tree;

        // if errors occurred, log the error count
        if error_count != 0 {
            warn!(
                "Onedrive Driver: {} errors occurred while building the tree for drive {}.",
                error_count,
                state.my_drive_id.clone()
            );
        }

        let folder = folder.unwrap_folder();

        Ok(folder.into())
    }
}

// auth
const AUTH_URL: &str = "https://login.microsoftonline.com/common/oauth2/v2.0/token";

#[allow(unused)]
#[derive(Debug, Deserialize)]
/// The response json when request `AUTH_URL`.
struct AccessTokenResponse {
    access_token: String,
    token_type: String,
    expires_in: i64,
    scope: String,
    refresh_token: String,
}

async fn fetch_access_token(config: &OnedriveConfig) -> Result<AccessTokenResponse, String> {
    let client = reqwest::Client::new();
    let res = client
        .post(AUTH_URL)
        .form(&[
            ("client_id", &config.client_id),
            ("refresh_token", &config.refresh_token),
            ("requested_token_use", &"on_behalf_of".to_owned()),
            ("client_secret", &config.client_secret),
            ("grant_type", &"refresh_token".to_owned()),
        ])
        .send()
        .await;
    match res {
        Ok(res) => {
            if let Ok(body) = res.json::<AccessTokenResponse>().await {
                Ok(body)
            } else {
                Err("Failed to parse response".to_owned())
            }
        }
        Err(e) => Err(format!("Failed to request access token: {}", e)),
    }
}

// get my drive id
const MY_DRIVE_URL: &str = "https://graph.microsoft.com/v1.0/me/drive";

#[derive(Debug, Deserialize)]
/// Response json when request `MY_DRIVE_URL`.
struct MyDrive {
    id: String,
}

async fn get_my_od_id(access_token: &str) -> Result<String, String> {
    let client = reqwest::Client::new();
    let res = client
        .get(MY_DRIVE_URL)
        .header("Authorization", format!("Bearer {}", access_token))
        .send()
        .await;
    match res {
        Ok(res) => {
            if let Ok(body) = res.json::<MyDrive>().await {
                Ok(body.id)
            } else {
                Err("Failed to parse response".into())
            }
        }
        Err(e) => Err(format!("Failed to request my drive id: {}", e).into()),
    }
}

// recursive get all files
fn request_list_url(dir_id: &str, drive_id: &str) -> String {
    format!(
        "https://graph.microsoft.com/v1.0/drives/{}/items/{}/children?select=id,name,size,folder,lastModifiedDateTime,file,@microsoft.graph.downloadUrl",
        drive_id, dir_id
    )
}

#[derive(Debug, Deserialize)]
/// the file or folder item in the response json.
struct ResponseItem {
    id: String,
    name: String,
    size: i64,
    #[serde(rename = "@microsoft.graph.downloadUrl")]
    file_download_url: Option<String>,
    file: Option<ResponseFile>,
    folder: Option<ResponseFolder>,
    #[serde(rename = "lastModifiedDateTime")]
    last_modified_date_time: String,
}

#[derive(Debug, Deserialize)]
struct ResponseFolder {
    #[serde(rename = "childCount")]
    _child_count: i64,
}

#[derive(Debug, Deserialize)]
struct ResponseFile {
    #[allow(dead_code)]
    hashes: ResponseFileHashes,
    #[serde(rename = "mimeType")]
    _mime_type: String,
}

#[derive(Debug, Deserialize)]
struct ResponseFileHashes {
    #[serde(rename = "quickXorHash")]
    _quick_xor_hash: String,
}

#[derive(Debug, Deserialize)]
/// the response json when request the graphql api.
struct ResponseList {
    value: Vec<ResponseItem>,
}

#[derive(StaticDownloadLinkFile, VfsMeta, Clone, StaticCombinableFile)]
struct OnedriveFile {
    name: String,
    size: u64,
    last_modified: SystemTime,
    links: Vec<String>,
}

async fn request_list(drive_id: &str, dir_id: &str, token: &str) -> Result<ResponseList, String> {
    let client = reqwest::Client::new();
    let res = client
        .get(request_list_url(dir_id, drive_id))
        .header("Authorization", format!("Bearer {}", token))
        .send()
        .await;
    match res {
        Ok(res) => {
            let body = match res.json::<ResponseList>().await {
                Ok(body) => body,
                Err(_) => return Err("Failed to parse response".to_owned()),
            };
            Ok(body)
        }
        Err(_) => Err("Failed to request list".to_owned()),
    }
}

/// internal struct to represent a folder in onedrive
struct OneDriveFolder {
    id: String,
    name: String,
    size: i64,
    last_modified: SystemTime,
    children: Vec<OneDriveItem>,
}

impl Into<CombinableDir<StaticCombinableFile>> for OneDriveFolder {
    fn into(self) -> CombinableDir<StaticCombinableFile> {
        let name = self.name;
        let (files, folders) = divide_items(self.children);
        let folders = folders
            .into_iter()
            .map(|folder| folder.into())
            .collect::<Vec<_>>();
        CombinableDir::new(
            name,
            files.into_iter().map(|file| file.into()).collect(),
            folders,
        )
    }
}

/// internal enum to represent a file or a folder in onedrive
enum OneDriveItem {
    File(OnedriveFile),
    Folder(OneDriveFolder),
    Unknown,
}

impl OneDriveItem {
    #[allow(unused)]
    pub fn unwrap_file(self) -> OnedriveFile {
        match self {
            OneDriveItem::File(this) => this,
            _ => panic!("Trying unwrap to file but it isn't file."),
        }
    }

    pub fn unwrap_folder(self) -> OneDriveFolder {
        match self {
            OneDriveItem::Folder(this) => this,
            _ => panic!("Trying unwrap to folder but it isn't folder."),
        }
    }

    #[allow(unused)]
    pub fn is_unknown(&self) -> bool {
        match self {
            OneDriveItem::Unknown => true,
            _ => false,
        }
    }
}

impl Into<OneDriveItem> for ResponseItem {
    fn into(self) -> OneDriveItem {
        match (self.file, self.folder, self.file_download_url) {
            (Some(_), None, Some(url)) => OneDriveItem::File(OnedriveFile {
                name: self.name,
                size: self.size as u64,
                links: vec![url],
                last_modified: DateTime::<Utc>::from(
                    DateTime::parse_from_rfc3339(self.last_modified_date_time.as_str()).unwrap(),
                )
                .into(),
            }),
            (None, Some(_), None) => OneDriveItem::Folder(OneDriveFolder {
                id: self.id,
                name: self.name,
                size: self.size,
                children: Vec::new(),
                last_modified: DateTime::<Utc>::from(
                    DateTime::parse_from_rfc3339(self.last_modified_date_time.as_str()).unwrap(),
                )
                .into(),
            }),
            _ => OneDriveItem::Unknown,
        }
    }
}

fn divide_items(items: Vec<OneDriveItem>) -> (Vec<OnedriveFile>, Vec<OneDriveFolder>) {
    let items = items.into_iter().filter(|item| match item {
        OneDriveItem::File(_) => true,
        OneDriveItem::Folder(_) => true,
        OneDriveItem::Unknown => false,
    });
    let mut files = Vec::new();
    let mut folders = Vec::new();
    for item in items {
        match item {
            OneDriveItem::File(file) => files.push(file),
            OneDriveItem::Folder(folder) => folders.push(folder),
            OneDriveItem::Unknown => (),
        }
    }
    (files, folders)
}

impl Into<CombinableDir<OnedriveFile>> for OneDriveFolder {
    fn into(self) -> CombinableDir<OnedriveFile> {
        let name = self.name;
        let (files, folders) = divide_items(self.children);
        let folders = folders
            .into_iter()
            .map(|folder| folder.into())
            .collect::<Vec<_>>();
        CombinableDir::new(name, files, folders)
    }
}

type RequestTreeResult = (OneDriveItem, i64);

// 1st: root, 2nd: error count
fn build_tree<'a>(
    access_token: String,
    drive_id: String,
    dir_id: String,
    dir_name: String,
    size: i64,
    last_modified_time: SystemTime,
) -> Pin<Box<dyn Future<Output = RequestTreeResult> + 'static + Send>> {
    Box::pin(async move {
        // request the graphql api
        let res = request_list(drive_id.as_str(), dir_id.as_str(), access_token.as_str()).await;
        if res.is_err() {
            return (OneDriveItem::Unknown, 1);
        }
        let list = res.unwrap().value;

        // initial error count and two vectors to store files and folders
        let mut error_count = 0;

        // divide the items in the list into files and folders
        let onedrive_items = list.into_iter().map(|item| item.into()).collect::<Vec<_>>();
        let (files, folders) = divide_items(onedrive_items);

        // graphql api will only return the id of the folders, so we can build the tree recursively
        let folders = folders
            .into_iter()
            .map(|folder| {
                build_tree(
                    access_token.clone(),
                    drive_id.clone(),
                    folder.id.clone(),
                    folder.name.clone(),
                    folder.size,
                    folder.last_modified,
                )
            })
            .collect::<Vec<_>>();

        // wait for all the folders to be built
        let futures = folders
            .into_iter()
            .map(|f| tokio::spawn(f))
            .collect::<Vec<_>>();
        let mut folders = Vec::with_capacity(futures.len());
        for future in futures {
            let (folder, count) = future.await.unwrap();
            error_count += count;
            folders.push(folder);
        }

        let files = files
            .into_iter()
            .map(|i| OneDriveItem::File(i))
            .collect::<Vec<_>>();

        let children = files.into_iter().chain(folders.into_iter()).collect();

        // return the result
        (
            OneDriveItem::Folder(OneDriveFolder {
                id: dir_id,
                name: dir_name,
                size,
                last_modified: last_modified_time,
                children,
            }),
            error_count,
        )
    })
}

#[cfg(test)]
mod test {
    use super::*;


    /// To test this driver, a real onedrive account is needed.
    /// You should create a file named `test.config.json` in the root directory of this project.
    const TEST_CONFIG_PATH: &str = "test.config.json";
    fn read_test_config() -> OnedriveConfig {
        let config = std::fs::read_to_string(TEST_CONFIG_PATH).unwrap();
        serde_json::from_str::<OnedriveConfig>(config.as_str()).unwrap()
    }

    #[tokio::test]
    async fn test_load_config() {
        read_test_config();
    }

    #[tokio::test]
    async fn test_get_my_od_id() {
        let config = read_test_config();
        let access_token = fetch_access_token(&config).await.unwrap();
        let AccessTokenResponse {
            access_token,
            ..
        } = access_token;
        let my_drive_id = get_my_od_id(&access_token).await.unwrap();
        assert!(!my_drive_id.is_empty());
    }

    #[tokio::test]
    async fn test_request_list() {
        let config = read_test_config();
        let access_token = fetch_access_token(&config).await.unwrap();
        let AccessTokenResponse {
            access_token,
            ..
        } = access_token;
        let my_drive_id = get_my_od_id(&access_token).await.unwrap();
        let _list = request_list(&my_drive_id, "root", &access_token).await.unwrap();
    }

    #[tokio::test]
    async fn test_build_tree() {
        let config = read_test_config();
        let access_token = fetch_access_token(&config).await.unwrap();
        let AccessTokenResponse {
            access_token,
            ..
        } = access_token;
        let my_drive_id = get_my_od_id(&access_token).await.unwrap();
        let tree = build_tree(
            access_token,
            my_drive_id,
            "root".to_owned(),
            "root".to_owned(),
            0,
            SystemTime::now(),
        )
        .await;
        let (folder, error_count) = tree;
        let folder = folder.unwrap_folder();
        assert_eq!(error_count, 0);
        assert_eq!(folder.name, "root");
    }
}