use super::{Hmac, StreamCipher};
use futures::prelude::*;
use std::{pin::Pin, task::Context, task::Poll};
#[pin_project::pin_project]
pub struct EncoderMiddleware<S> {
cipher_state: StreamCipher,
hmac: Hmac,
#[pin]
raw_sink: S,
}
impl<S> EncoderMiddleware<S> {
pub fn new(raw: S, cipher: StreamCipher, hmac: Hmac) -> EncoderMiddleware<S> {
EncoderMiddleware {
cipher_state: cipher,
hmac,
raw_sink: raw,
}
}
}
impl<S> Sink<Vec<u8>> for EncoderMiddleware<S>
where
S: Sink<Vec<u8>>,
{
type Error = S::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
let this = self.project();
Sink::poll_ready(this.raw_sink, cx)
}
fn start_send(self: Pin<&mut Self>, mut data_buf: Vec<u8>) -> Result<(), Self::Error> {
let this = self.project();
this.cipher_state.encrypt(&mut data_buf[..]);
let signature = this.hmac.sign(&data_buf[..]);
data_buf.extend_from_slice(signature.as_ref());
Sink::start_send(this.raw_sink, data_buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
let this = self.project();
Sink::poll_flush(this.raw_sink, cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
let this = self.project();
Sink::poll_close(this.raw_sink, cx)
}
}
impl<S> Stream for EncoderMiddleware<S>
where
S: Stream,
{
type Item = S::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.project();
Stream::poll_next(this.raw_sink, cx)
}
}