Skip to main content

mega/
easy.rs

1mod reader;
2mod util;
3
4pub use self::reader::FileDownloadReader;
5pub use self::util::ArcError;
6use crate::Command;
7use crate::Error;
8use crate::FetchNodesResponse;
9use crate::FileKey;
10use crate::GetAttributesResponse;
11use crate::ResponseData;
12use std::future::Future;
13use std::pin::Pin;
14// use std::sync::Arc;
15// use std::sync::Mutex;
16use tokio::io::AsyncRead;
17use tokio_stream::StreamExt;
18use tokio_util::io::StreamReader;
19
20/// A client
21#[derive(Debug, Clone)]
22pub struct Client {
23    /// The low-level api client
24    pub client: crate::Client,
25    // /// Client state
26    // state: Arc<Mutex<State>>,
27}
28
29impl Client {
30    /// Make a new client
31    pub fn new() -> Self {
32        Self {
33            client: crate::Client::new(),
34            /*
35            state: Arc::new(Mutex::new(State {
36                buffered_commands: Vec::with_capacity(4),
37                buffered_tx: Vec::with_capacity(4),
38            })),
39            */
40        }
41    }
42
43    /*
44    /// Queue a command to be sent
45    fn queue_command(
46        &self,
47        command: Command,
48    ) -> tokio::sync::oneshot::Receiver<Result<ResponseData, Error>> {
49        let (tx, rx) = tokio::sync::oneshot::channel();
50        {
51            let mut state = self.state.lock().unwrap();
52            state.buffered_commands.push(command);
53            state.buffered_tx.push(tx);
54        }
55        rx
56    }
57
58    /// Send all buffered commands
59    pub fn send_commands(&self) {
60        let (commands, tx) = {
61            let mut state = self.state.lock().unwrap();
62            if state.buffered_commands.is_empty() {
63                return;
64            }
65
66            let mut commands = Vec::with_capacity(4);
67            std::mem::swap(&mut commands, &mut state.buffered_commands);
68
69            let mut tx = Vec::with_capacity(4);
70            std::mem::swap(&mut tx, &mut state.buffered_tx);
71
72            (commands, tx)
73        };
74
75        let self_clone = self.clone();
76        tokio::spawn(async move {
77            let response = self_clone
78                .client
79                .execute_commands(&commands, None)
80                .await
81                .map_err(ArcError::new);
82            match response {
83                Ok(mut response) => {
84                    for tx in tx.into_iter().rev() {
85                        // The low-level api client ensures that the number of returned responses matches the number of input commands.
86                        let response = response.pop().unwrap();
87                        let response = response.into_result().map_err(Error::from);
88                        let _ = tx.send(response).is_ok();
89                    }
90                }
91                Err(error) => {
92                    for tx in tx {
93                        let _ = tx.send(Err(Error::BatchSend(error.clone()))).is_ok();
94                    }
95                }
96            };
97        });
98    }
99    */
100
101    /// Get attributes for a file.
102    pub fn get_attributes(
103        &self,
104        builder: GetAttributesBuilder,
105    ) -> impl Future<Output = Result<GetAttributesResponse, Error>> {
106        let command = Command::GetAttributes {
107            public_node_id: builder.public_node_id,
108            node_id: builder.node_id,
109            include_download_url: if builder.include_download_url {
110                Some(1)
111            } else {
112                None
113            },
114        };
115
116        async move {
117            let commands = [command];
118
119            let response = self
120                .client
121                .execute_commands(&commands, builder.reference_node_id.as_deref());
122
123            let response = match response.await?.swap_remove(0).into_result()? {
124                ResponseData::GetAttributes(response) => response,
125                _ => {
126                    return Err(Error::UnexpectedResponseDataType);
127                }
128            };
129
130            Ok(response)
131        }
132    }
133
134    /// Get the nodes for a folder node.
135    ///
136    /// This bypasses the command buffering system as it is more efficient for Mega's servers to process this alone.
137    pub async fn fetch_nodes(
138        &self,
139        node_id: Option<&str>,
140        recursive: bool,
141    ) -> Result<FetchNodesResponse, Error> {
142        let command = Command::FetchNodes {
143            c: 1,
144            recursive: u8::from(recursive),
145        };
146        let mut response = self
147            .client
148            .execute_commands(std::slice::from_ref(&command), node_id)
149            .await?;
150
151        // The low-level api client ensures that the number of returned responses matches the number of input commands.
152        let response = response.pop().unwrap();
153        let response = response.into_result().map_err(Error::from)?;
154        let response = match response {
155            ResponseData::FetchNodes(response) => response,
156            _ => {
157                return Err(Error::UnexpectedResponseDataType);
158            }
159        };
160
161        Ok(response)
162    }
163
164    /// Download a file without verifying its integrity.
165    ///
166    /// # Returns
167    /// Returns a reader.
168    pub async fn download_file_no_verify(
169        &self,
170        file_key: &FileKey,
171        url: &str,
172    ) -> Result<FileDownloadReader<Pin<Box<dyn AsyncRead + Send + Sync>>>, Error> {
173        let response = self
174            .client
175            .client
176            .get(url)
177            .send()
178            .await?
179            .error_for_status()?;
180
181        let stream_reader = StreamReader::new(
182            response
183                .bytes_stream()
184                .map(|result| result.map_err(std::io::Error::other)),
185        );
186        let stream_reader =
187            Box::into_pin(Box::new(stream_reader) as Box<dyn AsyncRead + Send + Sync>);
188
189        let reader = FileDownloadReader::new(stream_reader, file_key, false);
190
191        Ok(reader)
192    }
193
194    /// Download a file and verify its integrity.
195    ///
196    /// # Returns
197    /// Returns a reader.
198    pub async fn download_file(
199        &self,
200        file_key: &FileKey,
201        url: &str,
202    ) -> Result<FileDownloadReader<Pin<Box<dyn AsyncRead + Send + Sync>>>, Error> {
203        let response = self
204            .client
205            .client
206            .get(url)
207            .send()
208            .await?
209            .error_for_status()?;
210
211        let stream_reader = StreamReader::new(
212            response
213                .bytes_stream()
214                .map(|result| result.map_err(std::io::Error::other)),
215        );
216        let stream_reader =
217            Box::into_pin(Box::new(stream_reader) as Box<dyn AsyncRead + Send + Sync>);
218
219        let reader = FileDownloadReader::new(stream_reader, file_key, true);
220
221        Ok(reader)
222    }
223}
224
225impl Default for Client {
226    fn default() -> Self {
227        Self::new()
228    }
229}
230
231/*
232/// The client state
233#[derive(Debug)]
234struct State {
235    buffered_commands: Vec<Command>,
236    buffered_tx: Vec<tokio::sync::oneshot::Sender<Result<ResponseData, Error>>>,
237}
238*/
239
240/// A builder for a get_attributes call.
241#[derive(Debug)]
242pub struct GetAttributesBuilder {
243    /// The public id of the node.
244    ///
245    /// Mutually exclusive with `node_id`.
246    pub public_node_id: Option<String>,
247    /// The node id.
248    ///
249    /// Mutually exclusive with `public_file_id`.
250    pub node_id: Option<String>,
251    /// Whether this should include the download url
252    pub include_download_url: bool,
253
254    /// The reference node id.
255    pub reference_node_id: Option<String>,
256}
257
258impl GetAttributesBuilder {
259    /// Make a new builder.
260    pub fn new() -> Self {
261        Self {
262            public_node_id: None,
263            node_id: None,
264            include_download_url: false,
265            reference_node_id: None,
266        }
267    }
268
269    /// Set the public file id.
270    ///
271    /// Mutually exclusive with `node_id`.
272    pub fn public_node_id(&mut self, value: impl Into<String>) -> &mut Self {
273        self.public_node_id = Some(value.into());
274        self
275    }
276
277    /// Set the node id.
278    ///
279    /// Mutually exclusive with `public_node_id`.
280    pub fn node_id(&mut self, value: impl Into<String>) -> &mut Self {
281        self.node_id = Some(value.into());
282        self
283    }
284
285    /// Set the include_download_url field.
286    pub fn include_download_url(&mut self, value: bool) -> &mut Self {
287        self.include_download_url = value;
288        self
289    }
290
291    /// Set the reference node id.
292    pub fn reference_node_id(&mut self, value: impl Into<String>) -> &mut Self {
293        self.reference_node_id = Some(value.into());
294        self
295    }
296}
297
298impl Default for GetAttributesBuilder {
299    fn default() -> Self {
300        Self::new()
301    }
302}
303
304#[cfg(test)]
305mod test {
306    use super::*;
307    use crate::FolderKey;
308    use crate::test::*;
309    use std::collections::HashMap;
310    use tokio::io::AsyncReadExt;
311
312    /*
313    #[tokio::test]
314    async fn get_attributes() {
315        let client = Client::new();
316        let get_attributes_1_future = client.get_attributes(TEST_FILE_ID, false);
317        let get_attributes_2_future = client.get_attributes(TEST_FILE_ID, true);
318
319        let attributes_1 = get_attributes_1_future
320            .await
321            .expect("failed to get attributes");
322        assert!(attributes_1.download_url.is_none());
323        let attributes_2 = get_attributes_2_future
324            .await
325            .expect("failed to get attributes");
326        let file_attributes = attributes_1
327            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
328            .expect("failed to decode attributes");
329        assert!(file_attributes.name == "Doxygen_docs.zip");
330        assert!(attributes_2.download_url.is_some());
331        let file_attributes = attributes_2
332            .decode_attributes(TEST_FILE_KEY_KEY_DECODED)
333            .expect("failed to decode attributes");
334        assert!(file_attributes.name == "Doxygen_docs.zip");
335    }
336    */
337
338    #[tokio::test]
339    async fn fetch_nodes() {
340        let folder_key = FolderKey(TEST_FOLDER_KEY_DECODED);
341
342        let client = Client::new();
343        let response = client
344            .fetch_nodes(Some(TEST_FOLDER_ID), true)
345            .await
346            .expect("Failed to fetch nodes");
347        let mut key_map = HashMap::new();
348        key_map.insert(response.nodes[0].id.clone(), folder_key);
349        assert!(response.nodes.len() == 3);
350        let file_attributes = response
351            .nodes
352            .iter()
353            .find(|file| file.id == "oLkVhYqA")
354            .expect("Failed to locate file")
355            .decode_attributes(&key_map)
356            .expect("Failed to decode attributes");
357        assert!(file_attributes.name == "test");
358
359        let file_attributes = response
360            .nodes
361            .iter()
362            .find(|file| file.id == "kalwUahb")
363            .expect("Failed to locate file")
364            .decode_attributes(&key_map)
365            .expect("Failed to decode attributes");
366        assert!(file_attributes.name == "test.txt");
367
368        let file_attributes = &response
369            .nodes
370            .iter()
371            .find(|file| file.id == "IGlBlD6K")
372            .expect("Failed to locate file")
373            .decode_attributes(&key_map)
374            .expect("Failed to decode attributes");
375        assert!(file_attributes.name == "testfolder");
376    }
377
378    #[tokio::test]
379    async fn download_file_no_verify() {
380        let file_key = FileKey {
381            key: TEST_FILE_KEY_KEY_DECODED,
382            iv: TEST_FILE_KEY_IV_DECODED,
383            meta_mac: TEST_FILE_META_MAC_DECODED,
384        };
385
386        let client = Client::new();
387        let mut builder = GetAttributesBuilder::new();
388        builder
389            .include_download_url(true)
390            .public_node_id(TEST_FILE_ID);
391        let attributes = client
392            .get_attributes(builder)
393            .await
394            .expect("failed to get attributes");
395        let url = attributes.download_url.expect("missing download url");
396        let mut reader = client
397            .download_file_no_verify(&file_key, url.as_str())
398            .await
399            .expect("failed to get download stream");
400        let mut file = Vec::with_capacity(1024 * 1024);
401        tokio::io::copy(&mut reader, &mut file)
402            .await
403            .expect("failed to copy");
404
405        assert!(file == TEST_FILE_BYTES);
406    }
407
408    #[tokio::test]
409    async fn download_file_verify() {
410        let file_key = FileKey {
411            key: TEST_FILE_KEY_KEY_DECODED,
412            iv: TEST_FILE_KEY_IV_DECODED,
413            meta_mac: TEST_FILE_META_MAC_DECODED,
414        };
415
416        let client = Client::new();
417        {
418            let mut builder = GetAttributesBuilder::new();
419            builder
420                .include_download_url(true)
421                .public_node_id(TEST_FILE_ID);
422            let attributes = client
423                .get_attributes(builder)
424                .await
425                .expect("failed to get attributes");
426            let url = attributes.download_url.expect("missing download url");
427            let mut reader = client
428                .download_file(&file_key, url.as_str())
429                .await
430                .expect("failed to get download stream");
431            let mut file = Vec::with_capacity(1024 * 1024);
432            tokio::io::copy(&mut reader, &mut file)
433                .await
434                .expect("failed to copy");
435
436            assert!(file == TEST_FILE_BYTES);
437        }
438
439        {
440            let mut builder = GetAttributesBuilder::new();
441            builder
442                .include_download_url(true)
443                .public_node_id(TEST_FILE_ID);
444            let attributes = client
445                .get_attributes(builder)
446                .await
447                .expect("failed to get attributes");
448            let url = attributes.download_url.expect("missing download url");
449            let mut reader = client
450                .download_file(&file_key, url.as_str())
451                .await
452                .expect("failed to get download stream");
453            let mut file = Vec::new();
454            reader.read_to_end(&mut file).await.unwrap();
455
456            assert!(file == TEST_FILE_BYTES);
457        }
458    }
459}