1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use crate::errors::*;
use crate::{Mark, HEADER};
use bczhc_lib::fs::new_unique_file;
use bczhc_lib::io::ReadText;
use bczhc_lib::{rw_read, rw_write};
use byteorder::{BigEndian, ReadBytesExt};
use cfg_if::cfg_if;
use clap::ArgMatches;
use once_cell::sync::Lazy;
use std::ffi::OsStr;
use std::fs::File;
use std::io::Read;
use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::RwLock;
use tar::Archive;

static ARGUMENTS: Lazy<RwLock<Option<Arguments>>> = Lazy::new(|| RwLock::new(None));

pub fn main(matches: &ArgMatches) -> Result<()> {
    let output_dir = matches.value_of("output-dir").unwrap();

    rw_write!(ARGUMENTS).replace(Arguments {
        output_dir: String::from(output_dir),
    });

    let port = matches
        .value_of("port")
        .unwrap()
        .parse::<u16>()
        .map_err(|_| Error::InvalidPort)?;

    if matches.is_present("qr-code") {
        crate::qr::print_addr_qr(port)?;
    }

    let listener = TcpListener::bind(format!("0.0.0.0:{}", port))?;

    println!("Listening on port {}", port);

    let accept = listener.accept()?;
    println!("Accepted connection from {}", accept.1);

    let mut stream = accept.0;
    stream.check_header()?;
    let mark = stream.read_mark()?;

    match mark {
        Mark::File => receive_file(&mut stream),
        Mark::Text => receive_text(&mut stream),
        Mark::Tar => receive_files(&mut stream),
    }
}

fn receive_file<R>(stream: &mut R) -> Result<()>
where
    R: Read,
{
    let guard = rw_read!(ARGUMENTS);
    let output_dir = &guard.as_ref().unwrap().output_dir;

    let filename_length = stream.read_u32::<BigEndian>()? as usize;
    let mut filename_bytes = vec![0_u8; filename_length];
    stream.read_exact(&mut filename_bytes)?;

    #[allow(clippy::needless_late_init)]
    let filename;
    cfg_if! {
        if #[cfg(unix)] {
            use std::os::unix::prelude::OsStrExt;
            filename = OsStr::from_bytes(&filename_bytes);
        } else {
            filename =
                OsStr::new(std::str::from_utf8(&filename_bytes).map_err(|_| Error::InvalidCharset)?);
        }
    }

    let path = PathBuf::from(output_dir).join(filename);
    let path = new_unique_file(path)?;

    let mut file = File::options().write(true).truncate(true).open(path)?;

    std::io::copy(stream, &mut file)?;

    Ok(())
}

fn receive_files<R>(stream: &mut R) -> Result<()>
where
    R: Read,
{
    let guard = rw_read!(ARGUMENTS);
    let output_dir = &guard.as_ref().unwrap().output_dir;

    let mut archive = Archive::new(stream);

    archive.unpack(output_dir)?;

    Ok(())
}

fn receive_text<R>(stream: &mut R) -> Result<()>
where
    R: Read,
{
    let text = stream.read_text()?;
    println!("Received text:");
    println!("{}", text);
    Ok(())
}

trait Ext
where
    Self: Read,
{
    fn check_header(&mut self) -> Result<()> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        if &buf != HEADER {
            return Err(Error::InvalidHeader);
        }
        Ok(())
    }

    fn read_mark(&mut self) -> Result<Mark> {
        let mark = self.read_u8()?;
        num_traits::FromPrimitive::from_u8(mark).map_or_else(|| Err(Error::InvalidMark(mark)), Ok)
    }
}

struct Arguments {
    output_dir: String,
}

impl<R: Read> Ext for R {}