<commit_before>extern crate tiny_http;

use tiny_http::{Server, Response};
use std::net::{TcpStream, Shutdown};
use std::io::{copy, Write, Read};
use std::thread::{spawn, sleep};
use std::time::Duration;
use std::sync::mpsc::channel;
use std::sync::Arc;
use std::ops::Deref;

/// Stream that produces bytes very slowly
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct SlowByteSrc { val: u8, len: usize }
impl<'b> Read for SlowByteSrc {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        sleep(Duration::from_millis(100));
        let l = self.len.min(buf.len()).min(1000);
        buf[..l].fill(self.val);
        self.len -= l;
        Ok(l)
    }
}

/// crude impl of http `Transfer-Encoding: chunked`
fn encode_chunked(data: &mut dyn Read, output: &mut dyn Write) {
    let mut buf = [0u8; 4096];
    loop {
        let l = data.read(&mut buf).unwrap();
        write!(output, "{:X}\r\n", l).unwrap();
        output.write_all(&buf[..l]).unwrap();
        write!(output, "\r\n").unwrap();
        if l == 0 { break; }
    }
}

mod prompt_pipelining {
    use super::*;

    /// Check that pipelined requests on the same connection are received promptly.
    fn assert_requests_parsed_promptly(
        req_cnt: usize,
        req_body: &'static [u8],
        timeout: Duration,
        req_writer: impl FnOnce(&mut dyn Write) + Send + 'static
    ) {
        let resp_body = SlowByteSrc { val: 42, len: 1000_000 }; // very slow response body

        let server = Server::http("0.0.0.0:0").unwrap();
        let mut client = TcpStream::connect(server.server_addr()).unwrap();
        let (svr_send, svr_rcv) = channel();

        spawn(move || {
            for _ in 0..req_cnt {
                let mut req = server.recv().unwrap();
                // read the whole body of the request
                let mut body = Vec::new();
                req.as_reader().read_to_end(&mut body).unwrap();
                assert_eq!(req_body, body.as_slice());
                // The next pipelined request should now be available for parsing,
                // while we send the (possibly slow) response in another thread
                spawn(move || req.respond(
                    Response::empty(200).with_data(resp_body, Some(resp_body.len))
                ));
            }
            svr_send.send(()).unwrap();
        });

        spawn(move || req_writer(&mut client));

        // requests must be sent and received quickly (before timeout expires)
        svr_rcv.recv_timeout(timeout)
            .expect("Server did not finish reading pipelined requests quickly enough");
    }

    // short (but not trivial) request body
    const REQ_BODY: &'static [u8] = &[65; 5000];

    #[test]
    fn content_length() {
        assert_requests_parsed_promptly(5, REQ_BODY, Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n").unwrap();
                write!(wr, "Content-Length: {}\r\n\r\n", REQ_BODY.len()).unwrap();
                wr.write_all(REQ_BODY).unwrap();
            }
        });
    }

    #[test]
    fn chunked() {
        assert_requests_parsed_promptly(5, REQ_BODY, Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n").unwrap();
                write!(wr, "Transfer-Encoding: chunked\r\n\r\n").unwrap();
                let mut body = REQ_BODY;
                encode_chunked(&mut body, wr);
            }
        });
    }
}

mod prompt_responses {
    use super::*;

    /// Check that response is sent promptly without waiting for full request body.
    fn assert_responds_promptly(
        timeout: Duration,
        req_writer: impl FnOnce(&mut dyn Write) + Send + 'static
    ) {
        let server = Server::http("0.0.0.0:0").unwrap();
        let client = TcpStream::connect(server.server_addr()).unwrap();

        spawn(move || loop {
            // server attempts to respond immediately
            let req = server.recv().unwrap();
            req.respond(Response::empty(400)).unwrap();
        });

        let client = Arc::new(client);
        let client_write = Arc::clone(&client);
        // request written (possibly very slowly) in another thread
        spawn(move || req_writer(&mut client_write.deref()));

        // response should arrive quickly (before timeout expires)
        client.set_read_timeout(Some(timeout)).unwrap();
        let resp = client.deref().read(&mut [0u8; 4096]);
        client.shutdown(Shutdown::Both).unwrap();
        assert!(resp.is_ok(), "Server response was not sent promptly");
    }

    static SLOW_BODY: SlowByteSrc = SlowByteSrc { val: 65, len: 1000_000 };

    #[test]
    fn content_length_http11() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn content_length_http10() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.0\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn expect_continue() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Expect: 100 continue\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn chunked() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Transfer-Encoding: chunked\r\n\r\n").unwrap();
            encode_chunked(&mut SLOW_BODY.clone(), wr);
        });
    }
}<commit_msg>Additional tests for prompt handling of short requests.<commit_after>extern crate tiny_http;

use tiny_http::{Server, Response};
use std::net::{TcpStream, Shutdown};
use std::io::{copy, Write, Read};
use std::thread::{spawn, sleep};
use std::time::Duration;
use std::sync::mpsc::channel;
use std::sync::Arc;
use std::ops::Deref;

/// Stream that produces bytes very slowly
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct SlowByteSrc { val: u8, len: usize }
impl<'b> Read for SlowByteSrc {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        sleep(Duration::from_millis(100));
        let l = self.len.min(buf.len()).min(1000);
        buf[..l].fill(self.val);
        self.len -= l;
        Ok(l)
    }
}

/// crude impl of http `Transfer-Encoding: chunked`
fn encode_chunked(data: &mut dyn Read, output: &mut dyn Write) {
    let mut buf = [0u8; 4096];
    loop {
        let l = data.read(&mut buf).unwrap();
        write!(output, "{:X}\r\n", l).unwrap();
        output.write_all(&buf[..l]).unwrap();
        write!(output, "\r\n").unwrap();
        if l == 0 { break; }
    }
}

mod prompt_pipelining {
    use super::*;

    /// Check that pipelined requests on the same connection are received promptly.
    fn assert_requests_parsed_promptly(
        req_cnt: usize,
        req_body: &'static [u8],
        timeout: Duration,
        req_writer: impl FnOnce(&mut dyn Write) + Send + 'static
    ) {
        let resp_body = SlowByteSrc { val: 42, len: 1000_000 }; // very slow response body

        let server = Server::http("0.0.0.0:0").unwrap();
        let mut client = TcpStream::connect(server.server_addr()).unwrap();
        let (svr_send, svr_rcv) = channel();

        spawn(move || {
            for _ in 0..req_cnt {
                let mut req = server.recv().unwrap();
                // read the whole body of the request
                let mut body = Vec::new();
                req.as_reader().read_to_end(&mut body).unwrap();
                assert_eq!(req_body, body.as_slice());
                // The next pipelined request should now be available for parsing,
                // while we send the (possibly slow) response in another thread
                spawn(move || req.respond(
                    Response::empty(200).with_data(resp_body, Some(resp_body.len))
                ));
            }
            svr_send.send(()).unwrap();
        });

        spawn(move || req_writer(&mut client));

        // requests must be sent and received quickly (before timeout expires)
        svr_rcv.recv_timeout(timeout)
            .expect("Server did not finish reading pipelined requests quickly enough");
    }

    #[test]
    fn empty() {
        assert_requests_parsed_promptly(5, &[], Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n\r\n").unwrap();
            }
        });
    }

    #[test]
    fn content_length_short() {
        let body = &[65u8; 100]; // short but not trivial
        assert_requests_parsed_promptly(5, body, Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n").unwrap();
                write!(wr, "Content-Length: {}\r\n\r\n", body.len()).unwrap();
                wr.write_all(body).unwrap();
            }
        });
    }

    #[test]
    fn content_length_long() {
        let body = &[65u8; 10000]; // long enough that it won't be buffered
        assert_requests_parsed_promptly(5, body, Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n").unwrap();
                write!(wr, "Content-Length: {}\r\n\r\n", body.len()).unwrap();
                wr.write_all(body).unwrap();
            }
        });
    }

    #[test]
    fn chunked() {
        let body = &[65u8; 10000];
        assert_requests_parsed_promptly(5, body, Duration::from_millis(200), move |wr| {
            for _ in 0..5 {
                write!(wr, "GET / HTTP/1.1\r\n").unwrap();
                write!(wr, "Connection: keep-alive\r\n").unwrap();
                write!(wr, "Transfer-Encoding: chunked\r\n\r\n").unwrap();
                encode_chunked(&mut &body[..], wr);
            }
        });
    }
}

mod prompt_responses {
    use super::*;

    /// Check that response is sent promptly without waiting for full request body.
    fn assert_responds_promptly(
        timeout: Duration,
        req_writer: impl FnOnce(&mut dyn Write) + Send + 'static
    ) {
        let server = Server::http("0.0.0.0:0").unwrap();
        let client = TcpStream::connect(server.server_addr()).unwrap();

        spawn(move || loop {
            // server attempts to respond immediately
            let req = server.recv().unwrap();
            req.respond(Response::empty(400)).unwrap();
        });

        let client = Arc::new(client);
        let client_write = Arc::clone(&client);
        // request written (possibly very slowly) in another thread
        spawn(move || req_writer(&mut client_write.deref()));

        // response should arrive quickly (before timeout expires)
        client.set_read_timeout(Some(timeout)).unwrap();
        let resp = client.deref().read(&mut [0u8; 4096]);
        client.shutdown(Shutdown::Both).unwrap();
        assert!(resp.is_ok(), "Server response was not sent promptly");
    }

    static SLOW_BODY: SlowByteSrc = SlowByteSrc { val: 65, len: 1000_000 };

    #[test]
    fn content_length_http11() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn content_length_http10() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.0\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn expect_continue() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Expect: 100 continue\r\n").unwrap();
            write!(wr, "Content-Length: {}\r\n\r\n", SLOW_BODY.len).unwrap();
            copy(&mut SLOW_BODY.clone(), wr).unwrap();
        });
    }

    #[test]
    fn chunked() {
        assert_responds_promptly(Duration::from_millis(200), move |wr| {
            write!(wr, "GET / HTTP/1.1\r\n").unwrap();
            write!(wr, "Transfer-Encoding: chunked\r\n\r\n").unwrap();
            encode_chunked(&mut SLOW_BODY.clone(), wr);
        });
    }
}
{'fraction_non_alphanumeric': 0.1608741842464714, 'fraction_numerical': 0.018743360145697375, 'mean_word_length': 2.194909090909091, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5406673', 'n_tokens_mistral': 4360, 'n_tokens_neox': 4193, 'n_words': 926}