use bytes::Bytes;
use failure::Fail;
use futures::{Future, Stream};
use http::{request::Builder as RequestBuilder, response::Builder as ResponseBuilder};
use hyper::{
header::{HeaderValue, CONTENT_LENGTH},
Body, Chunk, Request, Response,
};
use std::borrow::Cow;
use crate::{Digest, ShaSize};
pub struct DigestBody {
inner: Chunk,
}
impl From<Chunk> for DigestBody {
fn from(b: Chunk) -> Self {
DigestBody { inner: b }
}
}
impl From<Bytes> for DigestBody {
fn from(b: Bytes) -> Self {
DigestBody { inner: b.into() }
}
}
impl From<Vec<u8>> for DigestBody {
fn from(b: Vec<u8>) -> Self {
DigestBody { inner: b.into() }
}
}
impl From<&'static [u8]> for DigestBody {
fn from(b: &'static [u8]) -> Self {
DigestBody { inner: b.into() }
}
}
impl From<Cow<'static, [u8]>> for DigestBody {
fn from(b: Cow<'static, [u8]>) -> Self {
DigestBody {
inner: b.into_owned().into(),
}
}
}
impl From<String> for DigestBody {
fn from(b: String) -> Self {
DigestBody { inner: b.into() }
}
}
impl From<&'static str> for DigestBody {
fn from(b: &'static str) -> Self {
DigestBody { inner: b.into() }
}
}
impl From<Cow<'static, str>> for DigestBody {
fn from(b: Cow<'static, str>) -> Self {
DigestBody {
inner: b.into_owned().into(),
}
}
}
pub trait WithDigest<T> {
fn with_digest<B: Into<DigestBody>>(&mut self, body: B, sha_size: ShaSize) -> http::Result<T>;
}
impl WithDigest<Response<Body>> for ResponseBuilder {
fn with_digest<B: Into<DigestBody>>(
&mut self,
body: B,
sha_size: ShaSize,
) -> http::Result<Response<Body>> {
let body: DigestBody = body.into();
let bytes = body.inner.into_bytes();
let digest = Digest::new(&bytes, sha_size);
self.header("Digest", HeaderValue::from_str(&digest.as_string())?)
.header(CONTENT_LENGTH, bytes.len())
.body(bytes.into())
}
}
impl WithDigest<Request<Body>> for RequestBuilder {
fn with_digest<B: Into<DigestBody>>(
&mut self,
body: B,
sha_size: ShaSize,
) -> http::Result<Request<Body>> {
let body: DigestBody = body.into();
let bytes = body.inner.into_bytes();
let digest = Digest::new(&bytes, sha_size);
self.header("Digest", HeaderValue::from_str(&digest.as_string())?)
.header(CONTENT_LENGTH, bytes.len())
.body(bytes.into())
}
}
#[derive(Debug, Fail)]
#[fail(display = "Error verifying digest")]
pub struct VerifyError;
pub trait VerifyDigest {
fn verify_digest(self) -> Box<dyn Future<Item = Self, Error = VerifyError> + Send>;
}
impl VerifyDigest for Request<Body> {
fn verify_digest(self) -> Box<dyn Future<Item = Request<Body>, Error = VerifyError> + Send> {
let maybe_digest = self.headers().get("Digest").cloned();
let (parts, body) = self.into_parts();
let f = verify(body, maybe_digest).map(|chunk| Request::from_parts(parts, chunk.into()));
Box::new(f)
}
}
impl VerifyDigest for Response<Body> {
fn verify_digest(self) -> Box<dyn Future<Item = Response<Body>, Error = VerifyError> + Send> {
let maybe_digest = self.headers().get("Digest").cloned();
let (parts, body) = self.into_parts();
let f = verify(body, maybe_digest).map(|chunk| Response::from_parts(parts, chunk.into()));
Box::new(f)
}
}
fn verify(
body: Body,
maybe_digest: Option<HeaderValue>,
) -> impl Future<Item = Chunk, Error = VerifyError> {
body.concat2()
.map_err(|_| VerifyError)
.and_then(move |body| {
maybe_digest
.ok_or(VerifyError)
.and_then(|digest| {
digest
.to_str()
.map(|s| s.to_owned())
.map_err(|_| VerifyError)
})
.and_then(|s| s.parse().map_err(|_| VerifyError))
.and_then(move |digest: Digest| {
digest
.verify(&body)
.map_err(|_| VerifyError)
.map(move |_| body)
})
})
}
#[cfg(test)]
mod tests {
use futures::Future;
use hyper::header::CONTENT_TYPE;
use hyper::{Request, Response};
use crate::{prelude::*, ShaSize};
#[test]
fn add_digest_to_request_256() {
add_digest_to_request(ShaSize::TwoFiftySix);
}
#[test]
fn add_digest_to_request_384() {
add_digest_to_request(ShaSize::ThreeEightyFour);
}
#[test]
fn add_digest_to_request_512() {
add_digest_to_request(ShaSize::FiveTwelve);
}
#[test]
fn add_digest_to_response_256() {
add_digest_to_response(ShaSize::TwoFiftySix);
}
#[test]
fn add_digest_to_response_384() {
add_digest_to_response(ShaSize::ThreeEightyFour);
}
#[test]
fn add_digest_to_response_512() {
add_digest_to_response(ShaSize::FiveTwelve);
}
fn add_digest_to_response(sha_size: ShaSize) {
let json = r#"{"Library":"Hyper"}"#;
let res = Response::builder()
.header(CONTENT_TYPE, "application/json")
.with_digest(json, sha_size)
.unwrap();
hyper::rt::run(
res.verify_digest()
.map(|_| ())
.map_err(|_| panic!("Failed to verify digest")),
);
}
fn add_digest_to_request(sha_size: ShaSize) {
let uri = "http://example.com";
let json = r#"{"Library":"Hyper"}"#;
let req = Request::post(uri)
.header(CONTENT_TYPE, "application/json")
.with_digest(json, sha_size)
.unwrap();
hyper::rt::run(
req.verify_digest()
.map(|_| ())
.map_err(|_| panic!("Failed to verify digest")),
);
}
}