Skip to main content

mega/
client.rs

1use crate::Command;
2use crate::Error;
3use crate::ErrorCode;
4use crate::Response;
5use crate::ResponseData;
6use std::sync::Arc;
7use std::sync::atomic::AtomicU64;
8use std::sync::atomic::Ordering;
9use std::time::Duration;
10use url::Url;
11
12/// A client
13#[derive(Debug, Clone)]
14pub struct Client {
15    /// The inner http client
16    pub client: reqwest::Client,
17
18    /// The sequence id
19    pub sequence_id: Arc<AtomicU64>,
20}
21
22impl Client {
23    /// Make a new client
24    pub fn new() -> Self {
25        Self {
26            client: reqwest::Client::new(),
27            sequence_id: Arc::new(AtomicU64::new(rand::random())),
28        }
29    }
30
31    /// Execute a series of commands.
32    ///
33    /// # Retries
34    /// If the client receives an EAGAIN,
35    /// it will attempt to retry the request.
36    /// After a number of tries with the same EAGAIN error,
37    /// the client will return EAGAIN to the caller.
38    pub async fn execute_commands(
39        &self,
40        commands: &[Command],
41        node: Option<&str>,
42    ) -> Result<Vec<Response<ResponseData>>, Error> {
43        const MAX_RETRIES: usize = 3;
44        const BASE_DELAY: u64 = 250;
45        const MAX_SEQUENCE_ID: u64 = 100_000;
46
47        let id = self.sequence_id.fetch_add(1, Ordering::Relaxed) % MAX_SEQUENCE_ID;
48        let mut url = Url::parse_with_params(
49            "https://g.api.mega.co.nz/cs",
50            [("id", itoa::Buffer::new().format(id))],
51        )?;
52        {
53            let mut query_pairs = url.query_pairs_mut();
54            if let Some(node) = node {
55                query_pairs.append_pair("n", node);
56            }
57        }
58
59        let mut retries = 0;
60        let response = loop {
61            let response: Response<Vec<_>> = self
62                .client
63                .post(url.as_str())
64                .json(commands)
65                .send()
66                .await?
67                .error_for_status()?
68                .json()
69                .await?;
70            let response = response.into_result();
71
72            if retries < MAX_RETRIES && matches!(response, Err(ErrorCode::EAGAIN)) {
73                let millis = BASE_DELAY * (1 << retries);
74                tokio::time::sleep(Duration::from_millis(millis)).await;
75                retries += 1;
76                continue;
77            }
78
79            break response;
80        };
81        let response = response?;
82
83        let commands_len = commands.len();
84        let response_len = response.len();
85        if response_len != commands_len {
86            return Err(Error::ResponseLengthMismatch {
87                expected: commands_len,
88                actual: response_len,
89            });
90        }
91
92        Ok(response)
93    }
94}
95
96impl Default for Client {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102#[cfg(test)]
103mod test {
104    use super::*;
105    use crate::test::*;
106    use crate::*;
107    use std::collections::HashMap;
108
109    #[tokio::test]
110    async fn execute_empty_commands() {
111        let client = Client::new();
112        let response = client
113            .execute_commands(&[], None)
114            .await
115            .expect("failed to execute commands");
116        assert!(response.is_empty());
117    }
118
119    #[tokio::test]
120    async fn execute_get_attributes_command() {
121        let client = Client::new();
122        let commands = vec![Command::GetAttributes {
123            public_node_id: Some(TEST_FILE_ID.into()),
124            node_id: None,
125            include_download_url: None,
126        }];
127        let mut response = client
128            .execute_commands(&commands, None)
129            .await
130            .expect("failed to execute commands");
131        assert!(response.len() == 1);
132        let response = response.swap_remove(0);
133        let response = response.into_result().expect("response was an error");
134        let response = match response {
135            ResponseData::GetAttributes(response) => response,
136            _ => panic!("unexpected response"),
137        };
138        assert!(response.download_url.is_none());
139        let file_attributes = response
140            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
141            .expect("failed to decode attributes");
142        assert!(file_attributes.name == "Doxygen_docs.zip");
143
144        let commands = vec![Command::GetAttributes {
145            public_node_id: Some(TEST_FILE_ID.into()),
146            node_id: None,
147            include_download_url: Some(1),
148        }];
149        let mut response = client
150            .execute_commands(&commands, None)
151            .await
152            .expect("failed to execute commands");
153        assert!(response.len() == 1);
154        let response = response.swap_remove(0);
155        let response = response.into_result().expect("response was an error");
156        let response = match response {
157            ResponseData::GetAttributes(response) => response,
158            _ => panic!("unexpected response"),
159        };
160        assert!(response.download_url.is_some());
161        let file_attributes = response
162            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
163            .expect("failed to decode attributes");
164        assert!(file_attributes.name == "Doxygen_docs.zip");
165    }
166
167    #[tokio::test]
168    async fn execute_fetch_nodes_command() {
169        let folder_key = FolderKey(TEST_FOLDER_KEY_DECODED);
170
171        let client = Client::new();
172        let commands = vec![Command::FetchNodes { c: 1, recursive: 1 }];
173        let mut response = client
174            .execute_commands(&commands, Some(TEST_FOLDER_ID))
175            .await
176            .expect("Failed to execute commands");
177        assert!(response.len() == 1);
178        let response = response.swap_remove(0);
179        let response = response.into_result().expect("Response was an error");
180        let response = match response {
181            ResponseData::FetchNodes(response) => response,
182            _ => panic!("Unexpected response"),
183        };
184        let mut key_map = HashMap::new();
185        key_map.insert(response.nodes[0].id.clone(), folder_key);
186        assert!(response.nodes.len() == 3);
187        let file_attributes = response
188            .nodes
189            .iter()
190            .find(|file| file.id == "oLkVhYqA")
191            .expect("failed to locate file")
192            .decode_attributes(&key_map)
193            .expect("failed to decode attributes");
194        assert!(file_attributes.name == "test");
195
196        let file_attributes = response
197            .nodes
198            .iter()
199            .find(|file| file.id == "kalwUahb")
200            .expect("failed to locate file")
201            .decode_attributes(&key_map)
202            .expect("failed to decode attributes");
203        assert!(file_attributes.name == "test.txt");
204
205        let file_attributes = &response
206            .nodes
207            .iter()
208            .find(|file| file.id == "IGlBlD6K")
209            .expect("failed to locate file")
210            .decode_attributes(&key_map)
211            .expect("failed to decode attributes");
212        assert!(file_attributes.name == "testfolder");
213    }
214}