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

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

use tower::Layer;

use super::{InstrumentOperation, MakeIdentity};

/// A [`Layer`] used to apply [`InstrumentOperation`].
#[derive(Debug)]
pub struct InstrumentLayer<RequestMakeFmt = MakeIdentity, ResponseMakeFmt = MakeIdentity> {
    operation_name: &'static str,
    make_request: RequestMakeFmt,
    make_response: ResponseMakeFmt,
}

impl InstrumentLayer {
    /// Constructs a new [`InstrumentLayer`] with no data redacted.
    pub fn new(operation_name: &'static str) -> Self {
        Self {
            operation_name,
            make_request: MakeIdentity,
            make_response: MakeIdentity,
        }
    }
}

impl<RequestMakeFmt, ResponseMakeFmt> InstrumentLayer<RequestMakeFmt, ResponseMakeFmt> {
    /// Configures the request format.
    ///
    /// The argument is typically [`RequestFmt`](super::sensitivity::RequestFmt).
    pub fn request_fmt<R>(self, make_request: R) -> InstrumentLayer<R, ResponseMakeFmt> {
        InstrumentLayer {
            operation_name: self.operation_name,
            make_request,
            make_response: self.make_response,
        }
    }

    /// Configures the response format.
    ///
    /// The argument is typically [`ResponseFmt`](super::sensitivity::ResponseFmt).
    pub fn response_fmt<R>(self, make_response: R) -> InstrumentLayer<RequestMakeFmt, R> {
        InstrumentLayer {
            operation_name: self.operation_name,
            make_request: self.make_request,
            make_response,
        }
    }
}

impl<S, RequestMakeFmt, ResponseMakeFmt> Layer<S> for InstrumentLayer<RequestMakeFmt, ResponseMakeFmt>
where
    RequestMakeFmt: Clone,
    ResponseMakeFmt: Clone,
{
    type Service = InstrumentOperation<S, RequestMakeFmt, ResponseMakeFmt>;

    fn layer(&self, service: S) -> Self::Service {
        InstrumentOperation::new(service, self.operation_name)
            .request_fmt(self.make_request.clone())
            .response_fmt(self.make_response.clone())
    }
}