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
108    #[tokio::test]
109    async fn execute_empty_commands() {
110        let client = Client::new();
111        let response = client
112            .execute_commands(&[], None)
113            .await
114            .expect("failed to execute commands");
115        assert!(response.is_empty());
116    }
117
118    #[tokio::test]
119    async fn execute_get_attributes_command() {
120        let client = Client::new();
121        let commands = vec![Command::GetAttributes {
122            public_node_id: Some(TEST_FILE_ID.into()),
123            node_id: None,
124            include_download_url: None,
125        }];
126        let mut response = client
127            .execute_commands(&commands, None)
128            .await
129            .expect("failed to execute commands");
130        assert!(response.len() == 1);
131        let response = response.swap_remove(0);
132        let response = response.into_result().expect("response was an error");
133        let response = match response {
134            ResponseData::GetAttributes(response) => response,
135            _ => panic!("unexpected response"),
136        };
137        assert!(response.download_url.is_none());
138        let file_attributes = response
139            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
140            .expect("failed to decode attributes");
141        assert!(file_attributes.name == "Doxygen_docs.zip");
142
143        let commands = vec![Command::GetAttributes {
144            public_node_id: Some(TEST_FILE_ID.into()),
145            node_id: None,
146            include_download_url: Some(1),
147        }];
148        let mut response = client
149            .execute_commands(&commands, None)
150            .await
151            .expect("failed to execute commands");
152        assert!(response.len() == 1);
153        let response = response.swap_remove(0);
154        let response = response.into_result().expect("response was an error");
155        let response = match response {
156            ResponseData::GetAttributes(response) => response,
157            _ => panic!("unexpected response"),
158        };
159        assert!(response.download_url.is_some());
160        let file_attributes = response
161            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
162            .expect("failed to decode attributes");
163        assert!(file_attributes.name == "Doxygen_docs.zip");
164    }
165
166    #[tokio::test]
167    async fn execute_fetch_nodes_command() {
168        let folder_key = FolderKey(TEST_FOLDER_KEY_DECODED);
169
170        let client = Client::new();
171        let commands = vec![Command::FetchNodes { c: 1, recursive: 1 }];
172        let mut response = client
173            .execute_commands(&commands, Some(TEST_FOLDER_ID))
174            .await
175            .expect("failed to execute commands");
176        assert!(response.len() == 1);
177        let response = response.swap_remove(0);
178        let response = response.into_result().expect("response was an error");
179        let response = match response {
180            ResponseData::FetchNodes(response) => response,
181            _ => panic!("unexpected response"),
182        };
183        assert!(response.nodes.len() == 3);
184        let file_attributes = response
185            .nodes
186            .iter()
187            .find(|file| file.id == "oLkVhYqA")
188            .expect("failed to locate file")
189            .decode_attributes(&folder_key)
190            .expect("failed to decode attributes");
191        assert!(file_attributes.name == "test");
192
193        let file_attributes = response
194            .nodes
195            .iter()
196            .find(|file| file.id == "kalwUahb")
197            .expect("failed to locate file")
198            .decode_attributes(&folder_key)
199            .expect("failed to decode attributes");
200        assert!(file_attributes.name == "test.txt");
201
202        let file_attributes = &response
203            .nodes
204            .iter()
205            .find(|file| file.id == "IGlBlD6K")
206            .expect("failed to locate file")
207            .decode_attributes(&folder_key)
208            .expect("failed to decode attributes");
209        assert!(file_attributes.name == "testfolder");
210    }
211}