如何在转发流的同时从流中复制数据

Question

我正在使用 hyper 0.12 构建代理服务。当从上游服务器接收到响应主体时，我想尽快将其转发回客户端，和将内容保存在缓冲区中以供以后处理。

所以我需要一个函数：

需要 Stream (a hyper::Body，准确地说）
returns 与输入流

Stream

也 returns 某种Future<Item = Vec<u8>, Error = ...> 当输出流完全消耗时，用输入流的缓冲内容解析

我这辈子都想不出该怎么做。

我想我正在寻找的函数看起来像这样：

type BufferFuture = Box<Future<Item = Vec<u8>, Error = ()>>;
pub fn copy_body(body: hyper::Body) -> (hyper::Body, BufferFuture) {
    let body2 = ... // ???
    let buffer = body.fold(Vec::<u8>::new(), |mut buf, chunk| {
        buf.extend_from_slice(&chunk);
        // ...somehow send this chunk to body2 also?
    });
    (body2, buffer);
}

以下是我尝试过的方法，它一直有效，直到 send_data() 失败（显然）。

type BufferFuture = Box<Future<Item = Vec<u8>, Error = ()>>;
pub fn copy_body(body: hyper::Body) -> (hyper::Body, BufferFuture) {
    let (mut sender, body2) = hyper::Body::channel();
    let consume =
        body.map_err(|_| ()).fold(Vec::<u8>::new(), move |mut buf, chunk| {
            buf.extend_from_slice(&chunk);

            // What to do if this fails?
            if sender.send_data(chunk).is_err() {}
            Box::new(future::ok(buf))
        });

    (body2, Box::new(consume));
}

然而，有件事告诉我我走错了路。

我找到了Sink.fanout()，这似乎是我想要的，但我没有Sink，而且我不知道如何构造一个。 hyper::Body 实现了 Stream 但没有实现 Sink.

Answer 1

我最终做的是实现一种新型流来满足我的需要。这似乎是必要的，因为 hyper::Body 没有实现 Sink，hyper::Chunk 也没有实现 Clone（Sink.fanout() 需要），所以我不能使用任何现有的组合器。

首先是一个结构，其中包含我们需要的所有详细信息和追加新块的方法，以及通知缓冲区已完成的方法。

struct BodyClone<T> {
    body: T,
    buffer: Option<Vec<u8>>,
    sender: Option<futures::sync::oneshot::Sender<Vec<u8>>>,
}

impl BodyClone<hyper::Body> {
    fn flush(&mut self) {
        if let (Some(buffer), Some(sender)) = (self.buffer.take(), self.sender.take()) {
            if sender.send(buffer).is_err() {}
        }
    }

    fn push(&mut self, chunk: &hyper::Chunk) {
        use hyper::body::Payload;

        let length = if let Some(buffer) = self.buffer.as_mut() {
            buffer.extend_from_slice(chunk);
            buffer.len() as u64
        } else {
            0
        };

        if let Some(content_length) = self.body.content_length() {
            if length >= content_length {
                self.flush();
            }
        }
    }
}

然后我为这个结构实现了 Stream 特性。

impl Stream for BodyClone<hyper::Body> {
    type Item = hyper::Chunk;
    type Error = hyper::Error;

    fn poll(&mut self) -> futures::Poll<Option<Self::Item>, Self::Error> {
        match self.body.poll() {
            Ok(Async::Ready(Some(chunk))) => {
                self.push(&chunk);
                Ok(Async::Ready(Some(chunk)))
            }
            Ok(Async::Ready(None)) => {
                self.flush();
                Ok(Async::Ready(None))
            }
            other => other,
        }
    }
}

我终于可以在 hyper::Body 上定义扩展方法了：

pub type BufferFuture = Box<Future<Item = Vec<u8>, Error = ()> + Send>;

trait CloneBody {
    fn clone_body(self) -> (hyper::Body, BufferFuture);
}

impl CloneBody for hyper::Body {
    fn clone_body(self) -> (hyper::Body, BufferFuture) {
        let (sender, receiver) = futures::sync::oneshot::channel();

        let cloning_stream = BodyClone {
            body: self,
            buffer: Some(Vec::new()),
            sender: Some(sender),
        };

        (
            hyper::Body::wrap_stream(cloning_stream),
            Box::new(receiver.map_err(|_| ())),
        )
    }
}

可以这样使用：

let (body: hyper::Body, buffer: BufferFuture) = body.clone_body();

如何在转发流的同时从流中复制数据

How to copy data from a stream while also forwarding a stream

rust

hyper