fs_share_utils/
lib.rs

1mod addr;
2mod utils;
3
4use indicatif::{ProgressBar, ProgressState, ProgressStyle};
5
6pub use addr::{get_receiver_addr, get_sender_addr};
7use std::fmt::Write as FWrite;
8use std::fs;
9use std::io::{BufWriter, Cursor};
10use std::path::PathBuf;
11use std::str::FromStr;
12use std::{
13    fs::File,
14    io::{self, BufReader, Write},
15    net::TcpStream,
16    path::Path,
17};
18
19pub use colored::*;
20
21pub use utils::sha256;
22
23const SEND_BUFFER_SIZE: usize = 32 * 1024;
24const RECEIVE_BUFFER_SIZE: usize = 32 * 1024;
25
26impl ShareFs for TcpStream {}
27
28pub trait ShareFs: Sized + io::Write + io::Read {
29    fn send_eof(&mut self) -> io::Result<()> {
30        write!(self, ":00:")
31    }
32    fn send_msg<T: AsRef<str>>(&mut self, value: T) -> io::Result<()> {
33        let msg = value.as_ref();
34        write!(self, "msg:{}:", msg.as_bytes().len())?;
35        let mut cursor = Cursor::new(msg);
36        transfer_with_progress(&mut cursor, self, msg.len(), SEND_BUFFER_SIZE)?;
37        let mut v = Vec::with_capacity(4);
38        transfer_without_progress(self, &mut v, 4, 4)?;
39        let prefix = std::str::from_utf8(&v).unwrap_or_default();
40        match prefix {
41            ":ss:" => {}
42            _ => {
43                eprintln!(".....Falid to Send.....");
44                std::process::exit(1);
45            }
46        }
47        Ok(())
48    }
49    fn receive<W: io::Write>(&mut self, stdout: &mut W) -> io::Result<bool> {
50        let mut v: Vec<u8> = Vec::with_capacity(4);
51        transfer_without_progress(self, &mut v, 4, 4)?;
52        let prefix = std::str::from_utf8(&v).unwrap_or_default();
53        match prefix {
54            ":00:" => Ok(false),
55            "msg:" => {
56                let len = read_num(self)?;
57                transfer_with_progress(self, stdout, len, RECEIVE_BUFFER_SIZE)?;
58                stdout.write_all(b"\n")?;
59                self.write_all(b":ss:")?;
60                Ok(true)
61            }
62            "fff:" => {
63                let name_length = read_num(self)?;
64                let file_length = read_num(self)?;
65                let mut name: Vec<u8> = Vec::with_capacity(name_length);
66                transfer_without_progress(self, &mut name, name_length, 1024)?;
67                let name = std::str::from_utf8(&name).unwrap_or_default();
68                // let mut file_writer = create_new_file(name)?; //
69                let f_n = PathBuf::from_str(name).unwrap();
70                let f = fs::File::create_new(f_n.file_name().unwrap())?;
71                let mut file_writer = BufWriter::new(f);
72                writeln!(
73                    stdout,
74                    "File Receiving: {}, Size: {} bytes",
75                    name, file_length
76                )?;
77                transfer_with_progress(self, &mut file_writer, file_length, RECEIVE_BUFFER_SIZE)?;
78                file_writer.flush()?;
79                self.write_all(b":ss:")?;
80                Ok(true)
81            }
82            _ => unreachable!(),
83        }
84    }
85    fn send_info<T: AsRef<str>>(&mut self, _name: T) -> io::Result<()> {
86        todo!()
87    }
88    fn send_file<P: AsRef<Path>, W: io::Write>(
89        &mut self,
90        file: P,
91        stdout: &mut W,
92    ) -> io::Result<()> {
93        let f_name = file
94            .as_ref()
95            .display()
96            .to_string()
97            .replace("\"", "")
98            .replace("'", "");
99        let f = File::open(file)?;
100        let file_len = f.metadata()?.len();
101        writeln!(
102            stdout,
103            "Sending file: {}, size: {} bytes",
104            &f_name, file_len
105        )?;
106        write!(
107            self,
108            "fff:{}:{}:{}",
109            f_name.as_bytes().len(),
110            file_len,
111            f_name
112        )?;
113        let mut reader = BufReader::new(f);
114        transfer_with_progress(&mut reader, self, file_len as usize, SEND_BUFFER_SIZE)?;
115        let mut buffer: Vec<u8> = vec![0; 4];
116        self.read_exact(&mut buffer[0..4])?;
117        let prefix = std::str::from_utf8(&buffer[0..4]).unwrap_or_default();
118        match prefix {
119            ":ss:" => {}
120            _ => {
121                writeln!(stdout, "{}", ".....Falid to Send.....".bold().red())?;
122                std::process::exit(1);
123            }
124        }
125        Ok(())
126    }
127}
128
129pub fn transfer_without_progress<R: io::Read, W: io::Write>(
130    r: &mut R,
131    w: &mut W,
132    n: usize,
133    buf_size: usize,
134) -> io::Result<()> {
135    let mut remain = n;
136    let mut buffer = vec![0; std::cmp::min(n, buf_size)];
137
138    while remain > 0 {
139        let to_read = std::cmp::min(buffer.len(), remain);
140        let read_count = r.read(&mut buffer[..to_read])?;
141        if read_count == 0 {
142            break;
143        }
144        w.write_all(&buffer[..read_count])?;
145        remain -= read_count;
146    }
147    Ok(())
148}
149
150fn transfer_with_progress<R: io::Read, W: io::Write>(
151    r: &mut R,
152    w: &mut W,
153    n: usize,
154    buf_size: usize,
155) -> io::Result<()> {
156    let mut remain = n;
157    let mut buffer = vec![0; std::cmp::min(n, buf_size)];
158
159    let mut i = 0;
160    let pb = ProgressBar::new(n as u64);
161    pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
162        .unwrap()
163        .with_key("eta", |state: &ProgressState, w: &mut dyn FWrite| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
164        .progress_chars("#>-"));
165
166    while remain > 0 {
167        let to_read = std::cmp::min(buffer.len(), remain);
168        let read_count = r.read(&mut buffer[..to_read])?;
169        if read_count == 0 {
170            break;
171        }
172        w.write_all(&buffer[..read_count])?;
173        remain -= read_count;
174
175        i += read_count as u64;
176        pb.set_position(i);
177    }
178    pb.finish();
179    Ok(())
180}
181
182pub fn read_num<R: io::Read>(r: &mut R) -> io::Result<usize> {
183    let mut buffer = [0; 1];
184    let mut num = 0;
185    loop {
186        let c = r.read(&mut buffer)?;
187        if c == 0 || !buffer[0].is_ascii_digit() {
188            break;
189        }
190        if num > usize::MAX / 10 {
191            return Err(io::Error::new(
192                io::ErrorKind::InvalidData,
193                "Number too large",
194            ));
195        }
196        num = num * 10 + (buffer[0] - b'0') as usize;
197    }
198    Ok(num)
199}
200/*
201fn create_new_file<T: AsRef<Path>>(name: T) -> io::Result<BufWriter<fs::File>> {
202    #[cfg(target_os = "android")]
203    let download_dir = PathBuf::from("/sdcard/Download");
204    #[cfg(not(target_os = "android"))]
205    let download_dir = dirs::download_dir().unwrap();
206
207    let file_name = download_dir.join(name);
208    dbg!(&file_name);
209    let path = ensure_unique_file(file_name);
210    let file = fs::File::create_new(path)?;
211    Ok(io::BufWriter::new(file))
212}
213
214
215fn ensure_unique_file<P: AsRef<Path>>(path: P) -> PathBuf {
216    let original_path = path.as_ref();
217    if !original_path.exists() {
218        return original_path.to_path_buf();
219    }
220
221    let mut count = 1;
222    let mut unique_path = original_path.to_path_buf();
223
224    while unique_path.exists() {
225        unique_path = original_path.with_file_name(format!(
226            "{}_{}{}",
227            original_path
228                .file_stem()
229                .unwrap_or_default()
230                .to_string_lossy(),
231            count,
232            original_path
233                .extension()
234                .map(|ext| format!(".{}", ext.to_string_lossy()))
235                .unwrap_or_default()
236        ));
237        count += 1;
238    }
239
240    unique_path
241}
242*/
243#[cfg(test)]
244mod tests {
245    use std::io::{self, sink, Cursor, Read};
246
247    use crate::{read_num, transfer_with_progress, transfer_without_progress};
248    const BUFFER_SIZE: usize = 1024;
249
250    #[test]
251    fn test_read_num_valid() {
252        let data = b"123045:abc";
253        let mut cursor = Cursor::new(data);
254        let result = read_num(&mut cursor).unwrap();
255        assert_eq!(result, 123045);
256    }
257    #[test]
258    fn test_read_num_empty() {
259        let data = b"";
260        let mut cursor = Cursor::new(data);
261        let result = read_num(&mut cursor).unwrap();
262        assert_eq!(result, 0);
263    }
264    #[test]
265    fn test_read_num_with_leading_zeros() {
266        let data = b"000000000000000000000000000000123";
267        let mut cursor = Cursor::new(data);
268        let result = read_num(&mut cursor).unwrap();
269        assert_eq!(result, 123);
270    }
271    #[test]
272    fn test_read_num_very_large_number() {
273        let data = format!("00000000{}::::::", usize::MAX);
274        let mut cursor = Cursor::new(data);
275        let result = read_num(&mut cursor).unwrap();
276        assert_eq!(result, usize::MAX);
277    }
278    #[test]
279    fn test_read_num_too_large_number() {
280        let data = b"9999999999999999999999999999999999";
281        let mut cursor = Cursor::new(data);
282        let result = read_num(&mut cursor);
283        assert!(result.is_err());
284        assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData);
285    }
286    #[test]
287    fn test_read_num() {
288        let data = b"0000100xyz43210?1*";
289        let mut cursor = Cursor::new(data);
290        let result = read_num(&mut cursor).unwrap(); // x
291        assert_eq!(result, 100);
292        let result = read_num(&mut cursor).unwrap(); // y
293        assert_eq!(result, 0);
294        let result = read_num(&mut cursor).unwrap(); // z
295        assert_eq!(result, 0);
296        let result = read_num(&mut cursor).unwrap(); // ?
297        assert_eq!(result, 43210);
298        let result = read_num(&mut cursor).unwrap(); // *
299        assert_eq!(result, 1);
300    }
301
302    #[test]
303    fn test_read_n_bytes_exact() {
304        let data = b"1234567890";
305        let mut reader = Cursor::new(data);
306        let mut writer = Vec::new();
307
308        // Read exactly 10 bytes
309        transfer_without_progress(&mut reader, &mut writer, 10, BUFFER_SIZE).unwrap();
310        assert_eq!(writer, data);
311    }
312    #[test]
313    fn test_read_n_bytes_partial() {
314        let data = b"1234567890";
315        let mut reader = Cursor::new(data);
316        let mut writer = Vec::new();
317
318        // Read only 5 bytes out of the 10 available
319        transfer_without_progress(&mut reader, &mut writer, 5, BUFFER_SIZE).unwrap();
320        assert_eq!(writer, b"12345");
321    }
322    #[test]
323    fn test_read_n_bytes_large_n() {
324        let data = b"1234567890";
325        let mut reader = Cursor::new(data);
326        let mut writer = Vec::new();
327
328        // Attempt to read more bytes than available
329        transfer_without_progress(&mut reader, &mut writer, 20, BUFFER_SIZE).unwrap();
330        assert_eq!(writer, data); // Writer should contain all available bytes
331    }
332    #[test]
333    fn test_read_n_bytes_empty_input() {
334        let data = b"";
335        let mut reader = Cursor::new(data);
336        let mut writer = Vec::new();
337
338        // Reading from an empty source should not fail
339        transfer_without_progress(&mut reader, &mut writer, 10, BUFFER_SIZE).unwrap();
340        assert!(writer.is_empty()); // Writer should remain empty
341    }
342    #[test]
343    fn test_read_n_bytes_zero_bytes() {
344        let data = b"1234567890";
345        let mut reader = Cursor::new(data);
346        let mut writer = Vec::new();
347
348        // Reading 0 bytes should result in no data written
349        transfer_without_progress(&mut reader, &mut writer, 0, BUFFER_SIZE).unwrap();
350        assert!(writer.is_empty());
351    }
352    #[test]
353    fn test_read_n_bytes_sink() {
354        let data = b"1234567890";
355        let mut reader = Cursor::new(data);
356        let mut writer = sink(); // /dev/null equivalent
357
358        // Reading into a sink; ensures no errors occur
359        transfer_without_progress(&mut reader, &mut writer, 5, BUFFER_SIZE).unwrap();
360        // Nothing to verify in the writer since it's a sink
361    }
362    #[test]
363    fn test_read_n_bytes_sink_large_data() {
364        use std::io::repeat;
365
366        // Simulate a source of infinite data (a stream of the byte `b'x'`)
367        let data = repeat(b'x'); // Infinite data
368        let mut reader = data.take(5 * 1024 * 1024 * 1024); // Limit to 5 GB
369        let mut writer = sink(); // Write to a sink (discarding data)
370
371        // Read 1.5 GB of data and write it to the sink
372        transfer_with_progress(
373            &mut reader,
374            &mut writer,
375            3 * 1024 * 1024 * 1024 / 2,
376            BUFFER_SIZE,
377        )
378        .unwrap();
379
380        // Test completes if no error occurs during the read/write
381    }
382    #[test]
383    fn test_read_n_bytes() {
384        let data = String::from_iter((0..20).into_iter().map(|_| {
385            format!(
386                "{}{}{}",
387                String::from_iter('a'..='z'),
388                String::from_iter('A'..='Z'),
389                String::from_iter('0'..='9')
390            )
391        }));
392        let mut reader = Cursor::new(data); // length = ((2 * 26) + 10) * 20 == 1240
393        let mut writer = Vec::new();
394        transfer_without_progress(&mut reader, &mut writer, 1, BUFFER_SIZE).unwrap();
395        assert_eq!(writer, b"a");
396        let mut writer = Vec::new();
397        transfer_without_progress(&mut reader, &mut writer, 15, BUFFER_SIZE).unwrap();
398        assert_eq!(writer, b"bcdefghijklmnop");
399        transfer_without_progress(&mut reader, &mut writer, 5, BUFFER_SIZE).unwrap();
400        let mut writer = Vec::new();
401        transfer_without_progress(&mut reader, &mut writer, 15, BUFFER_SIZE).unwrap();
402        assert_eq!(writer, b"vwxyzABCDEFGHIJ");
403        transfer_without_progress(&mut reader, &mut writer, (1240 - 36) - 20, BUFFER_SIZE).unwrap();
404        let mut writer = Vec::new();
405        transfer_with_progress(&mut reader, &mut writer, 20, BUFFER_SIZE).unwrap();
406        assert_eq!(writer, b"QRSTUVWXYZ0123456789");
407    }
408}
fs_share_utils/lib.rs

fs_share_utils/
lib.rs
