rust-kernel / user / fetch.rs

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
136
//! `fetch [-o file] http://host[:port]/path` (alias `wget`): download a
//! web page or file with HTTP/1.0 and print it, or save it with `-o`.
//!
//! HTTP is just text over a TCP connection. We send
//!     GET /path HTTP/1.0\r\n
//!     Host: example.com\r\n
//!     \r\n
//! and the server answers with a status line ("HTTP/1.0 200 OK"), header
//! lines, an empty line, then the body. With HTTP/1.0 the server closes
//! the connection when the body is done, so "read until EOF" is the whole
//! protocol on our side. (Plain HTTP only: HTTPS would need TLS.)
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

use core::fmt::Write;

/// "http://host:port/path" -> (host, port, path)
fn parse_url(url: &str) -> Option<(&str, u16, &str)> {
    let rest = url.strip_prefix("http://")?;
    let (hostport, path) = match rest.find('/') {
        Some(i) => (&rest[..i], &rest[i..]),
        None => (rest, "/"),
    };
    let (host, port) = match hostport.split_once(':') {
        Some((h, p)) => (h, p.parse().ok()?),
        None => (hostport, 80),
    };
    if host.is_empty() { None } else { Some((host, port, path)) }
}

fn main() -> i64 {
    let (mut out_path, mut url) = (None, None);
    let mut i = 1;
    while let Some(a) = ulib::arg(i) {
        if a == "-o" || a == "-O" {
            i += 1;
            out_path = ulib::arg(i);
        } else {
            url = Some(a);
        }
        i += 1;
    }
    let Some((host, port, path)) = url.and_then(parse_url) else {
        eprintln!("usage: fetch [-o file] http://host[:port]/path");
        return 2;
    };

    // 1. Name -> address (DNS), 2. connect (TCP handshake).
    let ip = match ulib::resolve(host) {
        Ok(ip) => ip,
        Err(e) => {
            eprintln!("fetch: {}: {}", host, ulib::strerror(e));
            return 1;
        }
    };
    let start = ulib::now_us();
    let fd = ulib::socket(ulib::SOCK_STREAM, 0);
    let r = if fd < 0 { fd } else { ulib::connect(fd, ip, port) };
    if r < 0 {
        eprintln!("fetch: connect to {}:{}: {}", ip, port, ulib::strerror(r));
        return 1;
    }

    // 3. Send the request.
    let mut request = alloc::string::String::new();
    let _ = write!(request, "GET {} HTTP/1.0\r\nHost: {}\r\nUser-Agent: rustbelt-fetch/0.12\r\nConnection: close\r\n\r\n", path, host);
    if ulib::write_all(fd, request.as_bytes()) < 0 {
        eprintln!("fetch: send failed");
        return 1;
    }

    // 4. Read the response: headers first (up to the empty line)...
    let mut head = alloc::vec::Vec::new();
    let mut buf = [0u8; 4096];
    let body_start = loop {
        let n = ulib::read(fd, &mut buf);
        if n <= 0 {
            eprintln!("fetch: connection closed before the headers ended");
            return 1;
        }
        head.extend_from_slice(&buf[..n as usize]);
        if let Some(i) = head.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        if head.len() > 16 * 1024 {
            eprintln!("fetch: headers too long");
            return 1;
        }
    };
    let status_line = head.split(|&b| b == b'\r').next().unwrap_or(b"");
    let status_line = core::str::from_utf8(status_line).unwrap_or("?");
    let code: u32 = status_line.split(' ').nth(1).and_then(|c| c.parse().ok()).unwrap_or(0);

    // ...then the body, to stdout or the file, until the server closes.
    let out = match out_path {
        Some(p) => ulib::open(p, ulib::O_CREAT | ulib::O_TRUNC),
        None => 1,
    };
    if out < 0 {
        eprintln!("fetch: {}: {}", out_path.unwrap(), ulib::strerror(out));
        return 1;
    }
    let mut total = head.len() - body_start;
    ulib::write_all(out, &head[body_start..]);
    loop {
        let n = ulib::read(fd, &mut buf);
        if n < 0 {
            eprintln!("fetch: read: {}", ulib::strerror(n));
            break;
        }
        if n == 0 {
            break; // the server closed: that's the end of the body
        }
        if ulib::write_all(out, &buf[..n as usize]) < 0 {
            eprintln!("fetch: write failed");
            return 1;
        }
        total += n as usize;
    }
    ulib::close(fd);
    let ms = ((ulib::now_us() - start) / 1000).max(1);
    if let Some(p) = out_path {
        ulib::close(out);
        println!(
            "fetch: {}, saved {} bytes to {} in {} ms ({} KiB/s)",
            status_line, total, p, ms, total as u64 * 1000 / 1024 / ms
        );
    }
    if (200..300).contains(&code) { 0 } else { 1 }
}

extern crate alloc;