moros/usr/
socket.rs

1use crate::api::console::Style;
2use crate::api::fs::IO;
3use crate::api::io;
4use crate::api::process::ExitCode;
5use crate::api::syscall;
6use crate::sys::console;
7use crate::sys::fs::OpenFlag;
8use crate::sys::net::SocketStatus;
9use crate::usr;
10
11use alloc::format;
12use alloc::vec;
13use alloc::vec::Vec;
14use bit_field::BitField;
15use core::str::{self, FromStr};
16use smoltcp::wire::IpAddress;
17
18pub fn main(args: &[&str]) -> Result<(), ExitCode> {
19    let mut listen = false;
20    let mut verbose = false;
21    let mut read_only = false;
22    let args: Vec<&str> = args.iter().filter_map(|arg| match *arg {
23        "-l" | "--listen" => {
24            listen = true;
25            None
26        }
27        "-r" | "--read" => {
28            read_only = true;
29            None
30        }
31        "-v" | "--verbose" => {
32            verbose = true;
33            None
34        }
35        _ => Some(*arg),
36    }).collect();
37
38    if verbose {
39        println!("MOROS Socket v0.2.0\n");
40    }
41
42    if args.len() != 2 {
43        help();
44        return Err(ExitCode::UsageError);
45    }
46    let (host, port) = match args[1].split_once(':') {
47        Some((h, p)) => (h, p),
48        None => ("0.0.0.0", args[1]),
49    };
50    let port: u16 = match port.parse() {
51        Ok(n) => n,
52        Err(_) => {
53            eprint!("Could not parse port");
54            return Err(ExitCode::UsageError);
55        }
56    };
57    let addr = if host.ends_with(char::is_numeric) {
58        IpAddress::from_str(host).expect("invalid address format")
59    } else {
60        match usr::host::resolve(host) {
61            Ok(ip_addr) => ip_addr,
62            Err(e) => {
63                error!("Could not resolve host: {:?}", e);
64                return Err(ExitCode::Failure);
65            }
66        }
67    };
68
69    let socket_path = "/dev/net/tcp";
70    let buf_len = if let Some(info) = syscall::info(socket_path) {
71        info.size() as usize
72    } else {
73        error!("Could not open '{}'", socket_path);
74        return Err(ExitCode::Failure);
75    };
76
77    let mut connected = false;
78    let stdin = 0;
79    let stdout = 1;
80    let flags = OpenFlag::Device as u8;
81    if let Some(handle) = syscall::open(socket_path, flags) {
82        if listen {
83            if syscall::listen(handle, port).is_err() {
84                error!("Could not listen to {}:{}", addr, port);
85                syscall::close(handle);
86                return Err(ExitCode::Failure);
87            }
88            if verbose {
89                debug!("Listening to {}:{}", addr, port);
90            }
91        } else {
92            if syscall::connect(handle, addr, port).is_ok() {
93                connected = true;
94            } else {
95                error!("Could not connect to {}:{}", addr, port);
96                syscall::close(handle);
97                return Err(ExitCode::Failure);
98            }
99            if verbose {
100                debug!("Connected to {}:{}", addr, port);
101            }
102        }
103
104        loop {
105            if console::end_of_text() || console::end_of_transmission() {
106                println!();
107                break;
108            }
109
110            if listen && !connected {
111                if syscall::accept(handle).is_ok() {
112                    connected = true;
113                } else {
114                    syscall::sleep(0.01);
115                    continue;
116                }
117            }
118
119            let list = vec![(stdin, IO::Read), (handle, IO::Read)];
120            if let Some((h, _)) = syscall::poll(&list) {
121                if h == stdin {
122                    let line = io::stdin().read_line().replace("\n", "\r\n");
123                    syscall::write(handle, line.as_bytes());
124                } else {
125                    let mut data = vec![0; buf_len];
126                    if let Some(bytes) = syscall::read(handle, &mut data) {
127                        data.resize(bytes, 0);
128                        syscall::write(stdout, &data);
129                    }
130                }
131            } else {
132                syscall::sleep(0.01);
133                if connected {
134                    let mut data = vec![0; 1]; // 1 byte status read
135                    match syscall::read(handle, &mut data) {
136                        Some(1) if is_closed(data[0]) => break,
137                        _ => continue,
138                    }
139                }
140            }
141        }
142        syscall::close(handle);
143        Ok(())
144    } else {
145        Err(ExitCode::Failure)
146    }
147}
148
149fn is_closed(status: u8) -> bool {
150    !status.get_bit(SocketStatus::MayRecv as usize)
151}
152
153fn help() {
154    let csi_option = Style::color("aqua");
155    let csi_title = Style::color("yellow");
156    let csi_reset = Style::reset();
157    println!(
158        "{}Usage:{} socket {}[<host>:]<port>{1}",
159        csi_title, csi_reset, csi_option
160    );
161    println!();
162    println!("{}Options:{}", csi_title, csi_reset);
163    println!(
164        "  {0}-l{1}, {0}--listen{1}    Listen to a local port",
165        csi_option, csi_reset
166    );
167    println!(
168        "  {0}-v{1}, {0}--verbose{1}   Increase verbosity",
169        csi_option, csi_reset
170    );
171    println!(
172        "  {0}-r{1}, {0}--read{1}      Read only connexion",
173        csi_option, csi_reset
174    );
175}