lmammino/digest.rs

## digest.rs
// api/shared/src/middleware/digest.rs

//! RFC 3230 Digest header middleware for AWS Lambda HTTP responses.
//!
//! This middleware computes the SHA-256 digest of lambda_http response bodies
//! and adds it as a `Digest` header in the format specified by RFC 3230.
//!
//! # Example
//!
//! ```rust,ignore
//! use shared::middleware::DigestLayer;
//! use tower::ServiceBuilder;
//!
//! let service = ServiceBuilder::new()
//!     .layer(DigestLayer)
//!     .service_fn(|event| handler(event));
//!
//! lambda_http::run(service).await
//! ```

use std::future::Future;
use std::ops::Deref;
use std::pin::Pin;
use std::task::{Context, Poll};

use base64::Engine;
use http::{HeaderValue, Response};
use lambda_http::Body;
use sha2::{Digest, Sha256};
use tower::{Layer, Service};

/// Tower Layer that adds RFC 3230 Digest headers to AWS Lambda HTTP responses.
///
/// This layer wraps services and computes SHA-256 digests of response bodies,
/// adding them as `Digest: sha-256=<base64>` headers.
#[derive(Clone, Copy, Debug, Default)]
pub struct DigestLayer;

impl DigestLayer {
    /// Creates a new `DigestLayer`.
    pub fn new() -> Self {
        Self
    }
}

impl<S> Layer<S> for DigestLayer {
    type Service = DigestService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        DigestService { inner }
    }
}

/// Tower Service that computes and attaches RFC 3230 Digest headers to Lambda responses.
///
/// This service wraps an inner service and intercepts its responses to:
/// 1. Extract the response body bytes
/// 2. Compute the SHA-256 digest
/// 3. Add the `Digest` header with the base64-encoded hash
#[derive(Clone, Debug)]
pub struct DigestService<S> {
    inner: S,
}

impl<S, ReqBody> Service<http::Request<ReqBody>> for DigestService<S>
where
    S: Service<http::Request<ReqBody>, Response = Response<Body>> + Send + 'static,
    S::Future: Send,
    S::Error: Send,
    ReqBody: Send + 'static,
{
    type Response = Response<Body>;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

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

    fn call(&mut self, request: http::Request<ReqBody>) -> Self::Future {
        let future = self.inner.call(request);
        Box::pin(async move {
            let response = future.await?;
            Ok(add_digest_header(response))
        })
    }
}

/// Computes the SHA-256 digest of the body and adds the Digest header.
fn add_digest_header(response: Response<Body>) -> Response<Body> {
    let (mut parts, body) = response.into_parts();

    // Get the body bytes - lambda_http::Body implements Deref<Target=[u8]>
    let body_bytes: &[u8] = body.deref();

    // Compute SHA-256 digest
    let mut hasher = Sha256::new();
    hasher.update(body_bytes);
    let hash = hasher.finalize();

    // Base64 encode the digest
    let digest_value = format!(
        "sha-256={}",
        base64::engine::general_purpose::STANDARD.encode(hash)
    );

    // Add the Digest header
    if let Ok(header_value) = HeaderValue::from_str(&digest_value) {
        parts.headers.insert("Digest", header_value);
    }

    // Reconstruct the response with the same body
    Response::from_parts(parts, body)
}

#[cfg(test)]
mod tests {
    use super::*;
    use http::{Request, StatusCode};
    use std::convert::Infallible;
    use tower::ServiceExt;

    /// A simple test service that returns a fixed body
    async fn echo_service(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
        Ok(Response::builder()
            .status(StatusCode::OK)
            .body(Body::from("hello world"))
            .unwrap())
    }

    #[tokio::test]
    async fn test_digest_header_added() {
        let service = DigestLayer::new().layer(tower::service_fn(echo_service));

        let request = Request::builder()
            .uri("/test")
            .body(Body::Empty)
            .unwrap();

        let response = service.oneshot(request).await.unwrap();

        // Check that the Digest header was added
        let digest_header = response.headers().get("Digest").unwrap();
        let digest_str = digest_header.to_str().unwrap();

        // Verify the format
        assert!(digest_str.starts_with("sha-256="));

        // Verify the actual digest value
        // SHA-256 of "hello world" = b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9
        // Base64 encoded = uU0nuZNNPgilLlLX2n2r+sSE7+N6U4DukIj3rOLvzek=
        assert_eq!(
            digest_str,
            "sha-256=uU0nuZNNPgilLlLX2n2r+sSE7+N6U4DukIj3rOLvzek="
        );
    }

    #[tokio::test]
    async fn test_empty_body_digest() {
        async fn empty_service(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
            Ok(Response::builder()
                .status(StatusCode::OK)
                .body(Body::Empty)
                .unwrap())
        }

        let service = DigestLayer::new().layer(tower::service_fn(empty_service));

        let request = Request::builder()
            .uri("/test")
            .body(Body::Empty)
            .unwrap();

        let response = service.oneshot(request).await.unwrap();

        let digest_header = response.headers().get("Digest").unwrap();
        let digest_str = digest_header.to_str().unwrap();

        // SHA-256 of empty string = e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
        // Base64 encoded = 47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=
        assert_eq!(
            digest_str,
            "sha-256=47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="
        );
    }

    #[tokio::test]
    async fn test_response_body_preserved() {
        let service = DigestLayer::new().layer(tower::service_fn(echo_service));

        let request = Request::builder()
            .uri("/test")
            .body(Body::Empty)
            .unwrap();

        let response = service.oneshot(request).await.unwrap();

        // Check that the body is preserved
        let body_bytes: &[u8] = response.body().deref();
        assert_eq!(body_bytes, b"hello world");
    }
}

## main.rs
// api/lambdas/some-lambda/src/main.rs

/// Usage example

mod http_handler;
use lambda_http::{run, tracing, Error};
use shared::middleware::DigestLayer;
use tower::ServiceBuilder;

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing::init_default_subscriber();

    // FROM HERE
    let service = ServiceBuilder::new()
        .layer(DigestLayer)
        .service_fn(http_handler::function_handler);

    run(service).await
}
	// api/shared/src/middleware/digest.rs

	//! RFC 3230 Digest header middleware for AWS Lambda HTTP responses.
	//!
	//! This middleware computes the SHA-256 digest of lambda_http response bodies
	//! and adds it as a `Digest` header in the format specified by RFC 3230.
	//!
	//! # Example
	//!
	//! ```rust,ignore
	//! use shared::middleware::DigestLayer;
	//! use tower::ServiceBuilder;
	//!
	//! let service = ServiceBuilder::new()
	//! .layer(DigestLayer)
	//! .service_fn(\|event\| handler(event));
	//!
	//! lambda_http::run(service).await
	//! ```

	use std::future::Future;
	use std::ops::Deref;
	use std::pin::Pin;
	use std::task::{Context, Poll};

	use base64::Engine;
	use http::{HeaderValue, Response};
	use lambda_http::Body;
	use sha2::{Digest, Sha256};
	use tower::{Layer, Service};

	/// Tower Layer that adds RFC 3230 Digest headers to AWS Lambda HTTP responses.
	///
	/// This layer wraps services and computes SHA-256 digests of response bodies,
	/// adding them as `Digest: sha-256=<base64>` headers.
	#[derive(Clone, Copy, Debug, Default)]
	pub struct DigestLayer;

	impl DigestLayer {
	/// Creates a new `DigestLayer`.
	pub fn new() -> Self {
	Self
	}
	}

	impl<S> Layer<S> for DigestLayer {
	type Service = DigestService<S>;

	fn layer(&self, inner: S) -> Self::Service {
	DigestService { inner }
	}
	}

	/// Tower Service that computes and attaches RFC 3230 Digest headers to Lambda responses.
	///
	/// This service wraps an inner service and intercepts its responses to:
	/// 1. Extract the response body bytes
	/// 2. Compute the SHA-256 digest
	/// 3. Add the `Digest` header with the base64-encoded hash
	#[derive(Clone, Debug)]
	pub struct DigestService<S> {
	inner: S,
	}

	impl<S, ReqBody> Service<http::Request<ReqBody>> for DigestService<S>
	where
	S: Service<http::Request<ReqBody>, Response = Response<Body>> + Send + 'static,
	S::Future: Send,
	S::Error: Send,
	ReqBody: Send + 'static,
	{
	type Response = Response<Body>;
	type Error = S::Error;
	type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

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

	fn call(&mut self, request: http::Request<ReqBody>) -> Self::Future {
	let future = self.inner.call(request);
	Box::pin(async move {
	let response = future.await?;
	Ok(add_digest_header(response))
	})
	}
	}

	/// Computes the SHA-256 digest of the body and adds the Digest header.
	fn add_digest_header(response: Response<Body>) -> Response<Body> {
	let (mut parts, body) = response.into_parts();

	// Get the body bytes - lambda_http::Body implements Deref<Target=[u8]>
	let body_bytes: &[u8] = body.deref();

	// Compute SHA-256 digest
	let mut hasher = Sha256::new();
	hasher.update(body_bytes);
	let hash = hasher.finalize();

	// Base64 encode the digest
	let digest_value = format!(
	"sha-256={}",
	base64::engine::general_purpose::STANDARD.encode(hash)
	);

	// Add the Digest header
	if let Ok(header_value) = HeaderValue::from_str(&digest_value) {
	parts.headers.insert("Digest", header_value);
	}

	// Reconstruct the response with the same body
	Response::from_parts(parts, body)
	}

	#[cfg(test)]
	mod tests {
	use super::*;
	use http::{Request, StatusCode};
	use std::convert::Infallible;
	use tower::ServiceExt;

	/// A simple test service that returns a fixed body
	async fn echo_service(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
	Ok(Response::builder()
	.status(StatusCode::OK)
	.body(Body::from("hello world"))
	.unwrap())
	}

	#[tokio::test]
	async fn test_digest_header_added() {
	let service = DigestLayer::new().layer(tower::service_fn(echo_service));

	let request = Request::builder()
	.uri("/test")
	.body(Body::Empty)
	.unwrap();

	let response = service.oneshot(request).await.unwrap();

	// Check that the Digest header was added
	let digest_header = response.headers().get("Digest").unwrap();
	let digest_str = digest_header.to_str().unwrap();

	// Verify the format
	assert!(digest_str.starts_with("sha-256="));

	// Verify the actual digest value
	// SHA-256 of "hello world" = b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9
	// Base64 encoded = uU0nuZNNPgilLlLX2n2r+sSE7+N6U4DukIj3rOLvzek=
	assert_eq!(
	digest_str,
	"sha-256=uU0nuZNNPgilLlLX2n2r+sSE7+N6U4DukIj3rOLvzek="
	);
	}

	#[tokio::test]
	async fn test_empty_body_digest() {
	async fn empty_service(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
	Ok(Response::builder()
	.status(StatusCode::OK)
	.body(Body::Empty)
	.unwrap())
	}

	let service = DigestLayer::new().layer(tower::service_fn(empty_service));

	let request = Request::builder()
	.uri("/test")
	.body(Body::Empty)
	.unwrap();

	let response = service.oneshot(request).await.unwrap();

	let digest_header = response.headers().get("Digest").unwrap();
	let digest_str = digest_header.to_str().unwrap();

	// SHA-256 of empty string = e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
	// Base64 encoded = 47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=
	assert_eq!(
	digest_str,
	"sha-256=47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU="
	);
	}

	#[tokio::test]
	async fn test_response_body_preserved() {
	let service = DigestLayer::new().layer(tower::service_fn(echo_service));

	let request = Request::builder()
	.uri("/test")
	.body(Body::Empty)
	.unwrap();

	let response = service.oneshot(request).await.unwrap();

	// Check that the body is preserved
	let body_bytes: &[u8] = response.body().deref();
	assert_eq!(body_bytes, b"hello world");
	}
	}
	// api/lambdas/some-lambda/src/main.rs

	/// Usage example

	mod http_handler;
	use lambda_http::{run, tracing, Error};
	use shared::middleware::DigestLayer;
	use tower::ServiceBuilder;

	#[tokio::main]
	async fn main() -> Result<(), Error> {
	tracing::init_default_subscriber();

	// FROM HERE
	let service = ServiceBuilder::new()
	.layer(DigestLayer)
	.service_fn(http_handler::function_handler);

	run(service).await
	}