When I use ByteStream, how should I add custom headers?

[lixiaofei123]

When I use ByteStream, how should I add headers? I'm trying to implement reverse proxy functionality with Rocket, allowing users to stream downloads by passing a URL parameter. My current code is as follows:

#[rocket::get("/proxy/<_..>")]
async fn handle_proxy(proxyreq: ProxyReq) -> Result<ByteStream![Bytes], (Status, String)> {
    let url = proxyreq.0;
    let host = proxyreq.1;
    if check_domain_is_allowed(&url, &host) {
        let proxy_resp = reqwest::get(&url)
            .await.map_err(|e| (Status::BadRequest, e.to_string()))?;
        let bytes_stream = proxy_resp.bytes_stream();
        Ok(ByteStream! {
            for await bytes in bytes_stream {
                match bytes {
                    Ok(bytes) => yield bytes,
                    Err(e) => {
                        eprintln!("error while streaming: {}", e);
                        break;
                    }
                }
            }
        })
    } else {
       Err((Status::BadRequest, "url is not allowed".to_owned()))
    }
}

It appears to be running fine so far. The only issue is that there is no Content-Length field in the response header, so the download progress cannot be displayed on the browser side. I want to retrieve the Content-Length field from the proxy_resp header and add it to the response header. However, I'm not sure how to do that. I tried using Responser, but without success.

[SergioBenitez]

A logical stream cannot have a length. This is a restriction of the HTTP spec, though not a limitation; logically, streams are not necessarily finite, so a length is not applicable. I believe you're using the wrong type for your use-case here, one governed by this restriction when such a thing isn't needed.

Everything in Rocket is streamed. The stream can have a known size or an unknown size. The ByteStream type is for creating "true" streams, in the sense that the stream may be endless. But in your case, you have an existing stream which can simply become the streamed body of a response.

Concretely, this means you can implement a simple Responder that forwards the stream. Something like:

struct ProxyStream(SomeStreamType);

impl<'r> Responder<'r, 'static> for ProxyStream {
    fn respond_to(self, _: &'r Request<'_>) -> response::Result<'static> {
    Response::build()
        .sized_body(some_length, self.0)
        .finalize();
    }
}

This assumes that SomeStreamType is AsyncRead, which it may or may not be. If it is instead a Stream, you can likely use StreamReader to convert it into an AsyncRead.

Then return a value of this type from your handler:

#[rocket::get("/proxy/<_..>")]
fn handle_proxy(proxyreq: ProxyReq) -> Result<ProxyStream, (Status, String)> {
    let (url, host) = proxyreq;
    if !check_domain_is_allowed(&url, &host) {
        return  Err((Status::BadRequest, "url is not allowed".to_owned()));
    }

    reqwest::get(&url)
        .await
        .map(ProxyStream)
        .map_err(|e| (Status::BadRequest, e.to_string()))
}

[lixiaofei123]

Thank you very much for your response. Based on your answer, I have updated my code. However, I am not very familiar with Rust, and I cannot determine the correct type of SomeStreamType in the following code.

struct ProxyReq(String, String);

#[rocket::async_trait]
impl<'a> FromRequest<'a> for ProxyReq {
    type Error = ();
    async fn from_request(request: &'a Request<'_>) -> request::Outcome<Self, ()> {
        let proxyurl = request
            .uri()
            .to_string()
            .strip_prefix("/proxy/")
            .unwrap_or("https://www.google.com")
            .to_owned();

        let mut host = "".to_owned();
        if let Some(host0) = request.host() {
            host = host0.domain().to_string();
        }

        Outcome::Success(ProxyReq(proxyurl, host))
    }
}

struct ProxyStream(SomeStreamType, Option<usize>);

impl<'r> Responder<'r, 'static> for ProxyStream {
    fn respond_to(self, _: &'r Request<'_>) -> response::Result<'static> {
        Response::build().sized_body(self.1, self.0).finalize();
    }
}

#[rocket::get("/proxy/<_..>")]
async fn handle_proxy(proxyreq: ProxyReq) -> Result<ProxyStream, (Status, String)> {
    let url = proxyreq.0;
    let host = proxyreq.1;
    if !check_domain_is_allowed(&url, &host) {
        return Err((Status::BadRequest, "url is not allowed".to_owned()));
    }

    let proxy_resp = reqwest::get(&url)
        .await
        .map_err(|e| (Status::BadRequest, e.to_string()))?;

    let mut content_length: Option<usize> = None;
    if let Some(length) = proxy_resp.headers().get("Content-Length"){
        let length = length.to_str()
            .map_err(|e| (Status::BadRequest, e.to_string()))?.parse::<usize>()
            .map_err(|e| (Status::BadRequest, e.to_string()))?;
        content_length = Some(length)
    }

    use rocket::futures::TryStreamExt;
    let bytes_stream = proxy_resp.bytes_stream().map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e));
    
    let reader = StreamReader::new(bytes_stream);
    Ok(ProxyStream(reader, content_length))
   
}

#[get("/")]
pub async fn index() -> Option<NamedFile> {
    NamedFile::open(Path::new("index.html")).await.ok()
}

use url::Url;

fn check_domain_is_allowed(proxyurl: &str, host: &str) -> bool {
    if let Ok(proxyurl) = Url::parse(proxyurl) {
        if let Some(domain) = proxyurl.domain() {
            return domain != host;
        }
    }
    false
}

#[launch]
fn rocket() -> _ {
    rocket::build().mount("/", routes![index, handle_proxy])
}