qs_core/
common.rs

1use std::path::Path;
2
3use async_compression::tokio::write::{GzipDecoder, GzipEncoder};
4use bincode::{Decode, Encode};
5use serde::{Deserialize, Serialize};
6use thiserror::Error;
7use tokio::io::AsyncWriteExt;
8
9/// Tree structure that represents the files that
10/// are being sent/received
11#[derive(Debug, PartialEq, Clone, Encode, Decode, Hash)]
12pub enum FileSendRecvTree {
13    File {
14        name: String,
15        skip: u64,
16        size: u64,
17    },
18    Dir {
19        name: String,
20        files: Vec<FileSendRecvTree>,
21    },
22}
23
24impl FileSendRecvTree {
25    /// Name of the file or directory
26    pub fn name(&self) -> &str {
27        match self {
28            FileSendRecvTree::File { name, .. } => name,
29            FileSendRecvTree::Dir { name, .. } => name,
30        }
31    }
32
33    /// Size of the tree in bytes
34    pub fn size(&self) -> u64 {
35        match self {
36            FileSendRecvTree::File { size, .. } => *size,
37            FileSendRecvTree::Dir { files, .. } => files.iter().map(|f| f.size()).sum(),
38        }
39    }
40
41    /// Number of bytes being partially skipped
42    /// [FileRecvSendTree] does not contain fully skipped files
43    pub fn skip(&self) -> u64 {
44        match self {
45            FileSendRecvTree::File { skip, .. } => *skip,
46            FileSendRecvTree::Dir { files, .. } => files.iter().map(|f| f.skip()).sum(),
47        }
48    }
49}
50
51/// Tree structure that represents the files that are available
52#[derive(Debug, PartialEq, Clone, Encode, Decode, Hash, Serialize, Deserialize)]
53pub enum FilesAvailable {
54    File {
55        name: String,
56        size: u64,
57    },
58    Dir {
59        name: String,
60        files: Vec<FilesAvailable>,
61    },
62}
63/// Get the available files
64pub fn get_files_available(path: &Path) -> std::io::Result<FilesAvailable> {
65    if path.is_file() {
66        Ok(FilesAvailable::File {
67            name: path.file_name().unwrap().to_str().unwrap().to_string(),
68            size: path.metadata()?.len(),
69        })
70    } else {
71        let mut files = Vec::new();
72        for entry in std::fs::read_dir(path)? {
73            let entry = entry?;
74            let path = entry.path();
75            files.push(get_files_available(&path)?);
76        }
77
78        Ok(FilesAvailable::Dir {
79            name: path.file_name().unwrap().to_str().unwrap().to_string(),
80            files,
81        })
82    }
83}
84
85impl FilesAvailable {
86    /// Name of the file or directory
87    pub fn name(&self) -> &str {
88        match self {
89            FilesAvailable::File { name, .. } => name,
90            FilesAvailable::Dir { name, .. } => name,
91        }
92    }
93
94    /// Size of the tree in bytes
95    pub fn size(&self) -> u64 {
96        match self {
97            FilesAvailable::File { size, .. } => *size,
98            FilesAvailable::Dir { files, .. } => files.iter().map(|f| f.size()).sum(),
99        }
100    }
101
102    /// Convert the tree to a [FileSendRecvTree]
103    pub fn to_send_recv_tree(&self) -> FileSendRecvTree {
104        match self {
105            FilesAvailable::File { name, size } => FileSendRecvTree::File {
106                name: name.to_string(),
107                skip: 0,
108                size: *size,
109            },
110            FilesAvailable::Dir { name, files } => FileSendRecvTree::Dir {
111                name: name.to_string(),
112                files: files.iter().map(|f| f.to_send_recv_tree()).collect(),
113            },
114        }
115    }
116
117    /// Fully/partially remove skipped files from the tree
118    /// - Returns [std::option::Option::None] if the tree is fully skipped
119    /// - panics if the tree roots do not match
120    pub fn remove_skipped(&self, to_skip: &FilesToSkip) -> Option<FileSendRecvTree> {
121        match (self, to_skip) {
122            (
123                FilesAvailable::File { name, size },
124                FilesToSkip::File {
125                    name: skip_name,
126                    skip,
127                },
128            ) => {
129                if name == skip_name && size <= skip {
130                    None
131                } else {
132                    Some(FileSendRecvTree::File {
133                        name: name.clone(),
134                        skip: *skip,
135                        size: *size,
136                    })
137                }
138            }
139            (
140                FilesAvailable::Dir { name, files },
141                FilesToSkip::Dir {
142                    name: skip_name,
143                    files: skip_files,
144                },
145            ) => {
146                if name != skip_name {
147                    panic!("Tree roots do not match");
148                }
149
150                let mut remaining_files = Vec::new();
151                for file in files {
152                    if let Some(skip_file) = skip_files.iter().find(|sf| match (file, sf) {
153                        (
154                            FilesAvailable::File { name, .. },
155                            FilesToSkip::File {
156                                name: skip_name, ..
157                            },
158                        ) => name == skip_name,
159                        (
160                            FilesAvailable::Dir { name, .. },
161                            FilesToSkip::Dir {
162                                name: skip_name, ..
163                            },
164                        ) => name == skip_name,
165                        _ => false,
166                    }) {
167                        if let Some(remaining) = file.remove_skipped(skip_file) {
168                            remaining_files.push(remaining);
169                        }
170                    } else {
171                        remaining_files.push(file.clone().to_send_recv_tree());
172                    }
173                }
174
175                if remaining_files.is_empty() {
176                    None
177                } else {
178                    Some(FileSendRecvTree::Dir {
179                        name: name.clone(),
180                        files: remaining_files,
181                    })
182                }
183            }
184            _ => panic!("Tree roots do not match"),
185        }
186    }
187
188    /// Compare two trees and return the files that can be skipped.
189    /// (e.g. compare local and remote files, returning those that can be skipped during transfer).
190    /// it is expected that ``self`` is larger than ``local_files``
191    /// # Returns
192    /// - [std::option::Option::None] if no files can be skipped
193    pub fn get_skippable(&self, local_files: &FilesAvailable) -> Option<FilesToSkip> {
194        match (self, local_files) {
195            (
196                FilesAvailable::File { name, .. },
197                FilesAvailable::File {
198                    name: local_name,
199                    size: local_size,
200                },
201            ) => {
202                if name == local_name {
203                    Some(FilesToSkip::File {
204                        name: name.clone(),
205                        skip: *local_size,
206                    })
207                } else {
208                    None
209                }
210            }
211            (
212                FilesAvailable::Dir { name, files },
213                FilesAvailable::Dir {
214                    name: local_name,
215                    files: local_files,
216                },
217            ) => {
218                if name != local_name {
219                    return None;
220                }
221
222                let mut skippable_files = Vec::new();
223                for file in files {
224                    if let Some(remote_file) = local_files.iter().find(|rf| match (file, rf) {
225                        (
226                            FilesAvailable::File { name, .. },
227                            FilesAvailable::File {
228                                name: local_name, ..
229                            },
230                        ) => name == local_name,
231                        (
232                            FilesAvailable::Dir { name, .. },
233                            FilesAvailable::Dir {
234                                name: local_name, ..
235                            },
236                        ) => name == local_name,
237                        _ => false,
238                    }) {
239                        if let Some(skippable) = file.get_skippable(remote_file) {
240                            skippable_files.push(skippable);
241                        }
242                    }
243                }
244
245                if skippable_files.is_empty() {
246                    None
247                } else {
248                    Some(FilesToSkip::Dir {
249                        name: name.clone(),
250                        files: skippable_files,
251                    })
252                }
253            }
254            _ => None,
255        }
256    }
257}
258
259/// Tree structure that represents files that have been requested for skipping
260#[derive(Debug, PartialEq, Clone, Encode, Decode, Hash)]
261pub enum FilesToSkip {
262    File {
263        name: String,
264        skip: u64,
265    },
266    Dir {
267        name: String,
268        files: Vec<FilesToSkip>,
269    },
270}
271
272impl FilesToSkip {
273    /// Name of the file or directory
274    pub fn name(&self) -> &str {
275        match self {
276            FilesToSkip::File { name, .. } => name,
277            FilesToSkip::Dir { name, .. } => name,
278        }
279    }
280
281    /// Number of bytes being skipped
282    /// This will include fully skipped files
283    pub fn skip(&self) -> u64 {
284        match self {
285            FilesToSkip::File { skip, .. } => *skip,
286            FilesToSkip::Dir { files, .. } => files.iter().map(|f| f.skip()).sum(),
287        }
288    }
289}
290
291pub async fn send_packet<P: Encode + std::fmt::Debug>(
292    packet: P,
293    conn: &iroh::endpoint::Connection,
294) -> std::io::Result<()> {
295    tracing::debug!("Sending packet: {:?}", packet);
296    let mut send = conn.open_uni().await?;
297
298    let data = bincode::encode_to_vec(&packet, bincode::config::standard()).unwrap();
299    let compressed = compress_gzip(&data).await?;
300    send.write_all(&compressed).await?;
301
302    send.flush().await?;
303    send.finish()?;
304
305    Ok(())
306}
307
308#[derive(Debug, Error)]
309pub enum PacketRecvError {
310    #[error("io error: {0}")]
311    IoError(#[from] std::io::Error),
312    #[error("encode error: {0}")]
313    EncodeError(#[from] bincode::error::DecodeError),
314    #[error("connection error: {0}")]
315    Connection(#[from] iroh::endpoint::ConnectionError),
316    #[error("read error {0}")]
317    Read(#[from] iroh::endpoint::ReadError),
318}
319
320pub async fn receive_packet<P: Decode + std::fmt::Debug>(
321    conn: &iroh::endpoint::Connection,
322) -> Result<P, PacketRecvError> {
323    let mut recv = conn.accept_uni().await?;
324    let mut buf = Vec::new();
325
326    loop {
327        let mut data = vec![0; 1024];
328        if let Some(n) = recv.read(&mut data).await? {
329            buf.extend_from_slice(&data[..n]);
330            continue;
331        }
332
333        break;
334    }
335
336    let decompressed = decompress_gzip(&buf).await?;
337
338    let packet = bincode::decode_from_slice(&decompressed, bincode::config::standard())?.0;
339
340    tracing::debug!("Received packet: {:?}", packet);
341
342    Ok(packet)
343}
344
345async fn compress_gzip(data: &[u8]) -> std::io::Result<Vec<u8>> {
346    let mut out = Vec::new();
347    let mut encoder = GzipEncoder::new(&mut out);
348    encoder.write_all(data).await?;
349    encoder.shutdown().await?;
350
351    Ok(out)
352}
353
354async fn decompress_gzip(data: &[u8]) -> std::io::Result<Vec<u8>> {
355    let mut out = Vec::new();
356    let mut decoder = GzipDecoder::new(&mut out);
357    decoder.write_all(data).await?;
358    decoder.shutdown().await?;
359
360    Ok(out)
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366    use pretty_assertions::assert_eq;
367
368    #[tokio::test]
369    async fn test_compression() {
370        let data = b"hellllllllllllllllllllllllo world";
371        let compressed = compress_gzip(data).await.unwrap();
372        let decompressed = decompress_gzip(&compressed).await.unwrap();
373
374        assert!(compressed.len() < data.len());
375        assert_eq!(data, &decompressed[..]);
376    }
377
378    #[test]
379    fn test_file_trees() {
380        let files_offered = FilesAvailable::Dir {
381            name: "root".to_string(),
382            files: vec![
383                FilesAvailable::File {
384                    name: "file1".to_string(),
385                    size: 10,
386                },
387                FilesAvailable::Dir {
388                    name: "dir1".to_string(),
389                    files: vec![
390                        FilesAvailable::File {
391                            name: "file2".to_string(),
392                            size: 20,
393                        },
394                        FilesAvailable::File {
395                            name: "file3".to_string(),
396                            size: 30,
397                        },
398                    ],
399                },
400            ],
401        };
402
403        let already_installed = FilesAvailable::Dir {
404            name: "root".to_string(),
405            files: vec![
406                FilesAvailable::File {
407                    name: "file1".to_string(),
408                    size: 10,
409                },
410                FilesAvailable::Dir {
411                    name: "dir1".to_string(),
412                    files: vec![FilesAvailable::File {
413                        name: "file2".to_string(),
414                        size: 15,
415                    }],
416                },
417            ],
418        };
419
420        let to_skip = files_offered.get_skippable(&already_installed).unwrap();
421        assert_eq!(
422            to_skip,
423            FilesToSkip::Dir {
424                name: "root".to_string(),
425                files: vec![
426                    FilesToSkip::File {
427                        name: "file1".to_string(),
428                        skip: 10
429                    },
430                    FilesToSkip::Dir {
431                        name: "dir1".to_string(),
432                        files: vec![FilesToSkip::File {
433                            name: "file2".to_string(),
434                            skip: 15
435                        }],
436                    },
437                ],
438            }
439        );
440
441        let new_tree_expected = FileSendRecvTree::Dir {
442            name: "root".to_string(),
443            files: vec![FileSendRecvTree::Dir {
444                name: "dir1".to_string(),
445                files: vec![
446                    FileSendRecvTree::File {
447                        name: "file2".to_string(),
448                        skip: 15,
449                        size: 20,
450                    },
451                    FileSendRecvTree::File {
452                        name: "file3".to_string(),
453                        skip: 0,
454                        size: 30,
455                    },
456                ],
457            }],
458        };
459
460        let new_tree = files_offered.remove_skipped(&to_skip).unwrap();
461        assert_eq!(new_tree, new_tree_expected);
462    }
463
464    #[test]
465    fn test_no_files_to_skip() {
466        let offered = FilesAvailable::Dir {
467            name: "root".to_string(),
468            files: vec![
469                FilesAvailable::File {
470                    name: "file1".to_string(),
471                    size: 10,
472                },
473                FilesAvailable::Dir {
474                    name: "dir1".to_string(),
475                    files: vec![
476                        FilesAvailable::File {
477                            name: "file2".to_string(),
478                            size: 20,
479                        },
480                        FilesAvailable::File {
481                            name: "file3".to_string(),
482                            size: 30,
483                        },
484                    ],
485                },
486            ],
487        };
488
489        let installed = FilesAvailable::Dir {
490            name: "root".to_string(),
491            files: vec![],
492        };
493
494        let to_skip = offered.get_skippable(&installed);
495        assert_eq!(to_skip, None);
496    }
497
498    #[test]
499    fn larger_directory() {
500        let offered = FilesAvailable::Dir {
501            name: "root".to_string(),
502            files: vec![
503                FilesAvailable::File {
504                    name: "file1".to_string(),
505                    size: 10,
506                },
507                FilesAvailable::Dir {
508                    name: "dir1".to_string(),
509                    files: vec![
510                        FilesAvailable::File {
511                            name: "file2".to_string(),
512                            size: 20,
513                        },
514                        FilesAvailable::File {
515                            name: "file3".to_string(),
516                            size: 30,
517                        },
518                        FilesAvailable::Dir {
519                            name: "dir2".to_string(),
520                            files: vec![FilesAvailable::File {
521                                name: "file4".to_string(),
522                                size: 40,
523                            }],
524                        },
525                    ],
526                },
527                FilesAvailable::Dir {
528                    name: "dir3".to_string(),
529                    files: vec![FilesAvailable::File {
530                        name: "file5".to_string(),
531                        size: 50,
532                    }],
533                },
534            ],
535        };
536
537        let installed = FilesAvailable::Dir {
538            name: "root".to_string(),
539            files: vec![
540                FilesAvailable::File {
541                    name: "file1".to_string(),
542                    size: 10,
543                },
544                FilesAvailable::Dir {
545                    name: "dir1".to_string(),
546                    files: vec![
547                        FilesAvailable::File {
548                            name: "file2".to_string(),
549                            size: 5,
550                        },
551                        FilesAvailable::Dir {
552                            name: "dir2".to_string(),
553                            files: vec![],
554                        },
555                    ],
556                },
557            ],
558        };
559
560        let to_skip = offered.get_skippable(&installed).unwrap();
561        assert_eq!(
562            to_skip,
563            FilesToSkip::Dir {
564                name: "root".to_string(),
565                files: vec![
566                    FilesToSkip::File {
567                        name: "file1".to_string(),
568                        skip: 10
569                    },
570                    FilesToSkip::Dir {
571                        name: "dir1".to_string(),
572                        files: vec![
573                            FilesToSkip::File {
574                                name: "file2".to_string(),
575                                skip: 5
576                            },
577                            // FilesToSkip::Dir {
578                            //     name: "dir2".to_string(),
579                            //     files: vec![],
580                            // },
581                        ],
582                    }
583                ]
584            }
585        );
586
587        let new_tree = offered.remove_skipped(&to_skip).unwrap();
588        let new_tree_expected = FileSendRecvTree::Dir {
589            name: "root".to_string(),
590            files: vec![
591                FileSendRecvTree::Dir {
592                    name: "dir1".to_string(),
593                    files: vec![
594                        FileSendRecvTree::File {
595                            name: "file2".to_string(),
596                            skip: 5,
597                            size: 20,
598                        },
599                        FileSendRecvTree::File {
600                            name: "file3".to_string(),
601                            skip: 0,
602                            size: 30,
603                        },
604                        FileSendRecvTree::Dir {
605                            name: "dir2".to_string(),
606                            files: vec![FileSendRecvTree::File {
607                                name: "file4".to_string(),
608                                skip: 0,
609                                size: 40,
610                            }],
611                        },
612                    ],
613                },
614                FileSendRecvTree::Dir {
615                    name: "dir3".to_string(),
616                    files: vec![FileSendRecvTree::File {
617                        name: "file5".to_string(),
618                        skip: 0,
619                        size: 50,
620                    }],
621                },
622            ],
623        };
624
625        assert_eq!(new_tree, new_tree_expected);
626    }
627}
qs_core/common.rs

qs_core/
common.rs
