changelog shortlog graph tags branches changeset files revisions annotate raw help

Mercurial > core / rust/lib/dl/tests/support/mod.rs

changeset 698: 96958d3eb5b0
parent: 966f92770ddf
author: Richard Westhaver <ellis@rwest.io>
date: Fri, 04 Oct 2024 22:04:59 -0400
permissions: -rw-r--r--
description: fixes 
     1 use std::{
     2   convert::Infallible,
     3   fs, io,
     4   net::SocketAddr,
     5   path::Path,
     6   sync::mpsc::{channel, Sender},
     7   thread,
     8 };
     9 
    10 use http_body_util::Full;
    11 use hyper::{body::Bytes, server::conn::http1, service::service_fn, Request};
    12 use tempfile::TempDir;
    13 
    14 pub fn tmp_dir() -> TempDir {
    15   tempfile::Builder::new()
    16     .prefix("rustup-download-test-")
    17     .tempdir()
    18     .expect("creating tempdir for test")
    19 }
    20 
    21 pub fn write_file(path: &Path, contents: &str) {
    22   let mut file = fs::OpenOptions::new()
    23     .write(true)
    24     .truncate(true)
    25     .create(true)
    26     .open(path)
    27     .expect("writing test data");
    28 
    29   io::Write::write_all(&mut file, contents.as_bytes())
    30     .expect("writing test data");
    31 
    32   file.sync_data().expect("writing test data");
    33 }
    34 
    35 // A dead simple hyper server implementation.
    36 // For more info, see:
    37 // https://hyper.rs/guides/1/server/hello-world/
    38 async fn run_server(
    39   addr_tx: Sender<SocketAddr>,
    40   addr: SocketAddr,
    41   contents: Vec<u8>,
    42 ) {
    43   let svc = service_fn(move |req: Request<hyper::body::Incoming>| {
    44     let contents = contents.clone();
    45     async move {
    46       let res = serve_contents(req, contents);
    47       Ok::<_, Infallible>(res)
    48     }
    49   });
    50 
    51   let listener = tokio::net::TcpListener::bind(&addr)
    52     .await
    53     .expect("can not bind");
    54 
    55   let addr = listener.local_addr().unwrap();
    56   addr_tx.send(addr).unwrap();
    57 
    58   loop {
    59     let (stream, _) = listener
    60       .accept()
    61       .await
    62       .expect("could not accept connection");
    63     let io = hyper_util::rt::TokioIo::new(stream);
    64 
    65     let svc = svc.clone();
    66     tokio::spawn(async move {
    67       if let Err(err) = http1::Builder::new().serve_connection(io, svc).await {
    68         eprintln!("failed to serve connection: {:?}", err);
    69       }
    70     });
    71   }
    72 }
    73 
    74 pub fn serve_file(contents: Vec<u8>) -> SocketAddr {
    75   let addr = ([127, 0, 0, 1], 0).into();
    76   let (addr_tx, addr_rx) = channel();
    77 
    78   thread::spawn(move || {
    79     let server = run_server(addr_tx, addr, contents);
    80     let rt =
    81       tokio::runtime::Runtime::new().expect("could not creating Runtime");
    82     rt.block_on(server);
    83   });
    84 
    85   let addr = addr_rx.recv();
    86   addr.unwrap()
    87 }
    88 
    89 fn serve_contents(
    90   req: hyper::Request<hyper::body::Incoming>,
    91   contents: Vec<u8>,
    92 ) -> hyper::Response<Full<Bytes>> {
    93   let mut range_header = None;
    94   let (status, body) =
    95     if let Some(range) = req.headers().get(hyper::header::RANGE) {
    96       // extract range "bytes={start}-"
    97       let range = range.to_str().expect("unexpected Range header");
    98       assert!(range.starts_with("bytes="));
    99       let range = range.trim_start_matches("bytes=");
   100       assert!(range.ends_with('-'));
   101       let range = range.trim_end_matches('-');
   102       assert_eq!(range.split('-').count(), 1);
   103       let start: u64 = range.parse().expect("unexpected Range header");
   104 
   105       range_header =
   106         Some(format!("bytes {}-{len}/{len}", start, len = contents.len()));
   107       (
   108         hyper::StatusCode::PARTIAL_CONTENT,
   109         contents[start as usize..].to_vec(),
   110       )
   111     } else {
   112       (hyper::StatusCode::OK, contents)
   113     };
   114 
   115   let mut res = hyper::Response::builder()
   116     .status(status)
   117     .header(hyper::header::CONTENT_LENGTH, body.len())
   118     .body(Full::new(Bytes::from(body)))
   119     .unwrap();
   120   if let Some(range) = range_header {
   121     res
   122       .headers_mut()
   123       .insert(hyper::header::CONTENT_RANGE, range.parse().unwrap());
   124   }
   125   res
   126 }