core/test/rust/runtime/http/streaming/upload.rs

use super::*;

fn chunk(socket: &mut BufReader<TcpStream>) -> Vec<u8> {
    let mut line = String::new();
    assert!(socket.read_line(&mut line).unwrap() > 0);
    let count = usize::from_str_radix(line.trim(), 16).unwrap();
    let mut body = vec![0; count];
    socket.read_exact(&mut body).unwrap();
    let mut ending = [0; 2];
    socket.read_exact(&mut ending).unwrap();
    assert_eq!(&ending, b"\r\n");
    body
}

fn request(body: AsyncStream<Event>) -> StreamRequest<Event> {
    StreamRequest {
        method: "POST".into(),
        target: "/stream".into(),
        headers: vec![],
        body,
    }
}

#[test]
fn duplex_upload_sends_first_chunk_before_source_eof_and_returns_response_early() {
    let (port, server) = peer(|listener| {
        let mut socket = accept(&listener);
        let head = headers(&mut socket).to_ascii_lowercase();
        assert!(head.contains("transfer-encoding: chunked"));
        assert!(!head.contains("content-length:"));
        assert_eq!(chunk(&mut socket), b"abc");
        socket
            .get_mut()
            .write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n3\r\nabc\r\n")
            .unwrap();
        assert_eq!(chunk(&mut socket), b"def");
        assert!(chunk(&mut socket).is_empty());
        socket
            .get_mut()
            .write_all(b"3\r\ndef\r\n0\r\n\r\n")
            .unwrap();
        assert_eq!(socket.read(&mut [0]).unwrap(), 0);
    });
    runtime()
        .run(async |scope| {
            let client = client(
                scope,
                port,
                &Limits {
                    body_bytes: 1,
                    ..Default::default()
                },
                &PoolLimits {
                    chunk_bytes: 3,
                    ..Default::default()
                },
            );
            let proceed = Arc::new(tokio::sync::Notify::new());
            let source = proceed.clone();
            let body = AsyncStream::new(futures_util::stream::unfold(0, move |index| {
                let source = source.clone();
                async move {
                    match index {
                        0 => Some((Ok(Ok(Bytes::from_static(b"abc"))), 1)),
                        1 => {
                            source.notified().await;
                            Some((Ok(Ok(Bytes::from_static(b"def"))), 2))
                        }
                        _ => None,
                    }
                }
            }));
            let alias = body.clone();
            let response = client
                .open_stream(request(body), event, event)
                .await
                .unwrap();
            assert_eq!(
                response.body.next().await.unwrap().unwrap().unwrap(),
                b"abc"[..]
            );
            proceed.notify_one();
            assert_eq!(collect(&response.body, 3).await, b"def");
            assert!(alias.next().await.unwrap().is_none());
            client.close().await.unwrap();
        })
        .unwrap();
    server.join().unwrap();
}

#[test]
fn upload_rejects_oversized_chunks_source_errors_and_stalled_reads() {
    for mode in 0..3 {
        let (port, server) = peer(|listener| {
            let mut socket = accept(&listener);
            let mut request = Vec::new();
            socket.read_to_end(&mut request).unwrap();
            assert!(!request.windows(8).any(|window| window == b"oversize"));
        });
        runtime()
            .run(async |scope| {
                let client = client(
                    scope,
                    port,
                    &Limits {
                        timeout_ms: 1000,
                        ..Default::default()
                    },
                    &PoolLimits {
                        chunk_bytes: 3,
                        read_ms: 40,
                        ..Default::default()
                    },
                );
                let body = match mode {
                    0 => AsyncStream::values([Ok(Bytes::from_static(b"oversize"))]),
                    1 => AsyncStream::values([Err("source failed".into())]),
                    _ => AsyncStream::new(futures_util::stream::pending()),
                };
                let alias = body.clone();
                assert!(
                    client
                        .open_stream(request(body), event, event)
                        .await
                        .is_err()
                );
                assert!(alias.next().await.unwrap().is_none());
                client.close().await.unwrap();
            })
            .unwrap();
        server.join().unwrap();
    }
}

#[test]
fn cancelling_request_before_headers_closes_upload_source_and_socket() {
    let (port, server) = peer(|listener| {
        let mut socket = accept(&listener);
        let mut bytes = vec![];
        socket.read_to_end(&mut bytes).unwrap();
    });
    runtime()
        .run(async |scope| {
            let client = client(scope, port, &Limits::default(), &PoolLimits::default());
            let body = AsyncStream::<Event>::new(futures_util::stream::pending());
            assert!(
                tokio::time::timeout(
                    Duration::from_millis(40),
                    client.open_stream(request(body.clone()), event, event)
                )
                .await
                .is_err()
            );
            tokio::time::timeout(Duration::from_secs(1), body.next())
                .await
                .unwrap()
                .unwrap();
            assert!(body.next().await.unwrap().is_none());
            client.close().await.unwrap();
        })
        .unwrap();
    server.join().unwrap();
}