1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

use crate::{error::BoxError, BoxFuture, Digest, Error, BUFFER_SIZE};
use bytes::{BufMut, Bytes, BytesMut};
use http::{header::HeaderName, HeaderValue, Request};
use http_body::{Body, Full};
use std::{
    error::Error as StdError,
    task::{Context, Poll},
};
use tower::{buffer::Buffer, Layer, Service};

/// Generate the digest header for this request
async fn sign_request<B>(
    digests: Vec<Digest>,
    overwrite: bool,
    req: Request<B>,
) -> Result<Request<Full<Bytes>>, Error>
where
    B: Body + Unpin,
    B::Error: Into<BoxError>,
{
    let (mut parts, mut req_body) = req.into_parts();

    // Read the whole body into memory (we have to do it at some point anyway)
    //
    // We could either enum dispatch or box the body to avoid having to read the body into memory if the value should not be overwritten and there's already a header present.
    // But alas, I can't be bothered right now.
    let mut body = BytesMut::new();
    while let Some(chunk) = req_body
        .data()
        .await
        .transpose()
        .map_err(|err| Error::BodyOperation(err.into()))?
    {
        body.put(chunk);
    }
    // We're done reading. Body can be frozen
    let body = body.freeze();

    let digest_header_name = HeaderName::from_static("digest");
    let digest_header_value = if overwrite || !parts.headers.contains_key(&digest_header_name) {
        // Generate hashes for all digests and construct the digest header value
        let mut digest_parts = Vec::new();
        for digest in digests {
            let name = digest.name();
            let value = digest.hash(&body)?.encode();

            digest_parts.push(format!("{name}={value}"));
        }

        HeaderValue::from_str(&digest_parts.join(","))?
    } else {
        parts.headers.get(&digest_header_name).cloned().unwrap()
    };

    parts
        .headers
        .insert(digest_header_name, digest_header_value);

    Ok(Request::from_parts(parts, Full::new(body)))
}

/// Tower layer for verifying the HTTP digest header
#[derive(Clone, Debug)]
pub struct SignerLayer {
    buffer_size: usize,
    digests: Vec<Digest>,
    overwrite: bool,
}

impl SignerLayer {
    /// Create a new signer layer with the default buffer size
    pub fn new(digests: Vec<Digest>, overwrite: bool) -> Self {
        Self::new_with_buffer_size(digests, BUFFER_SIZE, overwrite)
    }

    /// Create a new signer layer with a custom buffer size
    pub fn new_with_buffer_size(digests: Vec<Digest>, buffer_size: usize, overwrite: bool) -> Self {
        Self {
            buffer_size,
            digests,
            overwrite,
        }
    }
}

impl<S> Layer<S> for SignerLayer
where
    S: Service<Request<Full<Bytes>>> + Send + 'static,
    S::Error: StdError + Send + Sync,
    S::Future: Send,
{
    type Service = Signer<S>;

    fn layer(&self, service: S) -> Self::Service {
        Signer::new_with_buffer_size(
            service,
            self.buffer_size,
            self.digests.clone(),
            self.overwrite,
        )
    }
}

/// Service to generate the digest header for an HTTP request
#[derive(Clone)]
pub struct Signer<S>
where
    S: Service<Request<Full<Bytes>>>,
{
    inner: Buffer<S, Request<Full<Bytes>>>,
    digests: Vec<Digest>,
    overwrite: bool,
}

impl<S> Signer<S>
where
    S: Service<Request<Full<Bytes>>> + Send + 'static,
    S::Error: StdError + Send + Sync,
    S::Future: Send,
{
    /// Create a new signer service with the default buffer size
    pub fn new(inner: S, digests: Vec<Digest>, overwrite: bool) -> Self {
        Self::new_with_buffer_size(inner, BUFFER_SIZE, digests, overwrite)
    }

    /// Create a new signer service with a custom buffer size
    pub fn new_with_buffer_size(
        inner: S,
        buffer_size: usize,
        digests: Vec<Digest>,
        overwrite: bool,
    ) -> Self {
        let buffer = Buffer::new(inner, buffer_size);

        Self {
            inner: buffer,
            digests,
            overwrite,
        }
    }
}

// Not entirely sure which type bounds clippy thinks can be combined??
#[allow(clippy::type_repetition_in_bounds)]
impl<S, B> Service<Request<B>> for Signer<S>
where
    S: Service<Request<Full<Bytes>>> + Send + 'static,
    S::Error: Into<BoxError>,
    S::Future: Send,
    B: Body + Send + Unpin + 'static,
    B::Error: Into<BoxError>,
{
    type Response = S::Response;
    type Error = Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx).map_err(Error::InnerService)
    }

    fn call(&mut self, req: Request<B>) -> Self::Future {
        let (mut handle, digests, overwrite) =
            (self.inner.clone(), self.digests.clone(), self.overwrite);

        Box::pin(async move {
            let req = sign_request(digests, overwrite, req).await?;
            handle.call(req).await.map_err(Error::InnerService)
        })
    }
}