rama_http/layer/classify/
grpc_errors_as_failures.rs

1use super::{ClassifiedResponse, ClassifyEos, ClassifyResponse, SharedClassifier};
2use crate::{HeaderMap, Response};
3use bitflags::bitflags;
4use std::{fmt, num::NonZeroI32};
5
6/// gRPC status codes. Used in [`GrpcErrorsAsFailures`].
7///
8/// These variants match the [gRPC status codes].
9///
10/// [gRPC status codes]: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
11#[derive(Clone, Copy, Debug)]
12pub enum GrpcCode {
13    /// The operation completed successfully.
14    Ok,
15    /// The operation was cancelled.
16    Cancelled,
17    /// Unknown error.
18    Unknown,
19    /// Client specified an invalid argument.
20    InvalidArgument,
21    /// Deadline expired before operation could complete.
22    DeadlineExceeded,
23    /// Some requested entity was not found.
24    NotFound,
25    /// Some entity that we attempted to create already exists.
26    AlreadyExists,
27    /// The caller does not have permission to execute the specified operation.
28    PermissionDenied,
29    /// Some resource has been exhausted.
30    ResourceExhausted,
31    /// The system is not in a state required for the operation's execution.
32    FailedPrecondition,
33    /// The operation was aborted.
34    Aborted,
35    /// Operation was attempted past the valid range.
36    OutOfRange,
37    /// Operation is not implemented or not supported.
38    Unimplemented,
39    /// Internal error.
40    Internal,
41    /// The service is currently unavailable.
42    Unavailable,
43    /// Unrecoverable data loss or corruption.
44    DataLoss,
45    /// The request does not have valid authentication credentials
46    Unauthenticated,
47}
48
49impl GrpcCode {
50    pub(crate) const fn into_bitmask(self) -> GrpcCodeBitmask {
51        match self {
52            Self::Ok => GrpcCodeBitmask::OK,
53            Self::Cancelled => GrpcCodeBitmask::CANCELLED,
54            Self::Unknown => GrpcCodeBitmask::UNKNOWN,
55            Self::InvalidArgument => GrpcCodeBitmask::INVALID_ARGUMENT,
56            Self::DeadlineExceeded => GrpcCodeBitmask::DEADLINE_EXCEEDED,
57            Self::NotFound => GrpcCodeBitmask::NOT_FOUND,
58            Self::AlreadyExists => GrpcCodeBitmask::ALREADY_EXISTS,
59            Self::PermissionDenied => GrpcCodeBitmask::PERMISSION_DENIED,
60            Self::ResourceExhausted => GrpcCodeBitmask::RESOURCE_EXHAUSTED,
61            Self::FailedPrecondition => GrpcCodeBitmask::FAILED_PRECONDITION,
62            Self::Aborted => GrpcCodeBitmask::ABORTED,
63            Self::OutOfRange => GrpcCodeBitmask::OUT_OF_RANGE,
64            Self::Unimplemented => GrpcCodeBitmask::UNIMPLEMENTED,
65            Self::Internal => GrpcCodeBitmask::INTERNAL,
66            Self::Unavailable => GrpcCodeBitmask::UNAVAILABLE,
67            Self::DataLoss => GrpcCodeBitmask::DATA_LOSS,
68            Self::Unauthenticated => GrpcCodeBitmask::UNAUTHENTICATED,
69        }
70    }
71}
72
73bitflags! {
74    #[derive(Debug, Clone, Copy)]
75    pub(crate) struct GrpcCodeBitmask: u32 {
76        const OK                  = 0b00000000000000001;
77        const CANCELLED           = 0b00000000000000010;
78        const UNKNOWN             = 0b00000000000000100;
79        const INVALID_ARGUMENT    = 0b00000000000001000;
80        const DEADLINE_EXCEEDED   = 0b00000000000010000;
81        const NOT_FOUND           = 0b00000000000100000;
82        const ALREADY_EXISTS      = 0b00000000001000000;
83        const PERMISSION_DENIED   = 0b00000000010000000;
84        const RESOURCE_EXHAUSTED  = 0b00000000100000000;
85        const FAILED_PRECONDITION = 0b00000001000000000;
86        const ABORTED             = 0b00000010000000000;
87        const OUT_OF_RANGE        = 0b00000100000000000;
88        const UNIMPLEMENTED       = 0b00001000000000000;
89        const INTERNAL            = 0b00010000000000000;
90        const UNAVAILABLE         = 0b00100000000000000;
91        const DATA_LOSS           = 0b01000000000000000;
92        const UNAUTHENTICATED     = 0b10000000000000000;
93    }
94}
95
96impl GrpcCodeBitmask {
97    fn try_from_u32(code: u32) -> Option<Self> {
98        match code {
99            0 => Some(Self::OK),
100            1 => Some(Self::CANCELLED),
101            2 => Some(Self::UNKNOWN),
102            3 => Some(Self::INVALID_ARGUMENT),
103            4 => Some(Self::DEADLINE_EXCEEDED),
104            5 => Some(Self::NOT_FOUND),
105            6 => Some(Self::ALREADY_EXISTS),
106            7 => Some(Self::PERMISSION_DENIED),
107            8 => Some(Self::RESOURCE_EXHAUSTED),
108            9 => Some(Self::FAILED_PRECONDITION),
109            10 => Some(Self::ABORTED),
110            11 => Some(Self::OUT_OF_RANGE),
111            12 => Some(Self::UNIMPLEMENTED),
112            13 => Some(Self::INTERNAL),
113            14 => Some(Self::UNAVAILABLE),
114            15 => Some(Self::DATA_LOSS),
115            16 => Some(Self::UNAUTHENTICATED),
116            _ => None,
117        }
118    }
119}
120
121/// Response classifier for gRPC responses.
122///
123/// gRPC doesn't use normal HTTP statuses for indicating success or failure but instead a special
124/// header that might appear in a trailer.
125///
126/// Responses are considered successful if
127///
128/// - `grpc-status` header value contains a success value.
129/// - `grpc-status` header is missing.
130/// - `grpc-status` header value isn't a valid `String`.
131/// - `grpc-status` header value can't parsed into an `i32`.
132///
133/// All others are considered failures.
134#[derive(Debug, Clone)]
135pub struct GrpcErrorsAsFailures {
136    success_codes: GrpcCodeBitmask,
137}
138
139impl Default for GrpcErrorsAsFailures {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145impl GrpcErrorsAsFailures {
146    /// Create a new [`GrpcErrorsAsFailures`].
147    pub const fn new() -> Self {
148        Self {
149            success_codes: GrpcCodeBitmask::OK,
150        }
151    }
152
153    /// Change which gRPC codes are considered success.
154    ///
155    /// Defaults to only considering `Ok` as success.
156    ///
157    /// `Ok` will always be considered a success.
158    ///
159    /// # Example
160    ///
161    /// Servers might not want to consider `Invalid Argument` or `Not Found` as failures since
162    /// thats likely the clients fault:
163    ///
164    /// ```rust
165    /// use rama_http::layer::classify::{GrpcErrorsAsFailures, GrpcCode};
166    ///
167    /// let classifier = GrpcErrorsAsFailures::new()
168    ///     .with_success(GrpcCode::InvalidArgument)
169    ///     .with_success(GrpcCode::NotFound);
170    /// ```
171    pub fn with_success(mut self, code: GrpcCode) -> Self {
172        self.success_codes |= code.into_bitmask();
173        self
174    }
175
176    /// Change which gRPC codes are considered success.
177    ///
178    /// Same as [`Self::with_success`] but without consuming `self`.
179    pub fn set_success(&mut self, code: GrpcCode) -> &mut Self {
180        self.success_codes |= code.into_bitmask();
181        self
182    }
183
184    /// Returns a [`MakeClassifier`](super::MakeClassifier) that produces `GrpcErrorsAsFailures`.
185    ///
186    /// This is a convenience function that simply calls `SharedClassifier::new`.
187    pub fn make_classifier() -> SharedClassifier<Self> {
188        SharedClassifier::new(Self::new())
189    }
190}
191
192impl ClassifyResponse for GrpcErrorsAsFailures {
193    type FailureClass = GrpcFailureClass;
194    type ClassifyEos = GrpcEosErrorsAsFailures;
195
196    fn classify_response<B>(
197        self,
198        res: &Response<B>,
199    ) -> ClassifiedResponse<Self::FailureClass, Self::ClassifyEos> {
200        match classify_grpc_metadata(res.headers(), self.success_codes) {
201            ParsedGrpcStatus::Success
202            | ParsedGrpcStatus::HeaderNotString
203            | ParsedGrpcStatus::HeaderNotInt => ClassifiedResponse::Ready(Ok(())),
204            ParsedGrpcStatus::NonSuccess(status) => {
205                ClassifiedResponse::Ready(Err(GrpcFailureClass::Code(status)))
206            }
207            ParsedGrpcStatus::GrpcStatusHeaderMissing => {
208                ClassifiedResponse::RequiresEos(GrpcEosErrorsAsFailures {
209                    success_codes: self.success_codes,
210                })
211            }
212        }
213    }
214
215    fn classify_error<E>(self, error: &E) -> Self::FailureClass
216    where
217        E: fmt::Display,
218    {
219        GrpcFailureClass::Error(error.to_string())
220    }
221}
222
223/// The [`ClassifyEos`] for [`GrpcErrorsAsFailures`].
224#[derive(Debug, Clone)]
225pub struct GrpcEosErrorsAsFailures {
226    success_codes: GrpcCodeBitmask,
227}
228
229impl ClassifyEos for GrpcEosErrorsAsFailures {
230    type FailureClass = GrpcFailureClass;
231
232    fn classify_eos(self, trailers: Option<&HeaderMap>) -> Result<(), Self::FailureClass> {
233        if let Some(trailers) = trailers {
234            match classify_grpc_metadata(trailers, self.success_codes) {
235                ParsedGrpcStatus::Success
236                | ParsedGrpcStatus::GrpcStatusHeaderMissing
237                | ParsedGrpcStatus::HeaderNotString
238                | ParsedGrpcStatus::HeaderNotInt => Ok(()),
239                ParsedGrpcStatus::NonSuccess(status) => Err(GrpcFailureClass::Code(status)),
240            }
241        } else {
242            Ok(())
243        }
244    }
245
246    fn classify_error<E>(self, error: &E) -> Self::FailureClass
247    where
248        E: fmt::Display,
249    {
250        GrpcFailureClass::Error(error.to_string())
251    }
252}
253
254impl Default for GrpcEosErrorsAsFailures {
255    fn default() -> Self {
256        Self {
257            success_codes: GrpcCodeBitmask::OK,
258        }
259    }
260}
261
262/// The failure class for [`GrpcErrorsAsFailures`].
263#[derive(Debug)]
264pub enum GrpcFailureClass {
265    /// A gRPC response was classified as a failure with the corresponding status.
266    Code(std::num::NonZeroI32),
267    /// A gRPC response was classified as an error with the corresponding error description.
268    Error(String),
269}
270
271impl fmt::Display for GrpcFailureClass {
272    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
273        match self {
274            Self::Code(code) => write!(f, "Code: {}", code),
275            Self::Error(error) => write!(f, "Error: {}", error),
276        }
277    }
278}
279
280#[allow(clippy::match_result_ok)]
281pub(crate) fn classify_grpc_metadata(
282    headers: &HeaderMap,
283    success_codes: GrpcCodeBitmask,
284) -> ParsedGrpcStatus {
285    macro_rules! or_else {
286        ($expr:expr, $other:ident) => {
287            if let Some(value) = $expr {
288                value
289            } else {
290                return ParsedGrpcStatus::$other;
291            }
292        };
293    }
294
295    let status = or_else!(headers.get("grpc-status"), GrpcStatusHeaderMissing);
296    let status = or_else!(status.to_str().ok(), HeaderNotString);
297    let status = or_else!(status.parse::<i32>().ok(), HeaderNotInt);
298
299    if GrpcCodeBitmask::try_from_u32(status as _)
300        .filter(|code| success_codes.contains(*code))
301        .is_some()
302    {
303        ParsedGrpcStatus::Success
304    } else {
305        ParsedGrpcStatus::NonSuccess(NonZeroI32::new(status).unwrap())
306    }
307}
308
309#[derive(Debug, PartialEq, Eq, Hash)]
310pub(crate) enum ParsedGrpcStatus {
311    Success,
312    NonSuccess(NonZeroI32),
313    GrpcStatusHeaderMissing,
314    // these two are treated as `Success` but kept separate for clarity
315    HeaderNotString,
316    HeaderNotInt,
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    macro_rules! classify_grpc_metadata_test {
324        (
325            name: $name:ident,
326            status: $status:expr,
327            success_flags: $success_flags:expr,
328            expected: $expected:expr,
329        ) => {
330            #[test]
331            fn $name() {
332                let mut headers = HeaderMap::new();
333                headers.insert("grpc-status", $status.parse().unwrap());
334                let status = classify_grpc_metadata(&headers, $success_flags);
335                assert_eq!(status, $expected);
336            }
337        };
338    }
339
340    classify_grpc_metadata_test! {
341        name: basic_ok,
342        status: "0",
343        success_flags: GrpcCodeBitmask::OK,
344        expected: ParsedGrpcStatus::Success,
345    }
346
347    classify_grpc_metadata_test! {
348        name: basic_error,
349        status: "1",
350        success_flags: GrpcCodeBitmask::OK,
351        expected: ParsedGrpcStatus::NonSuccess(NonZeroI32::new(1).unwrap()),
352    }
353
354    classify_grpc_metadata_test! {
355        name: two_success_codes_first_matches,
356        status: "0",
357        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
358        expected: ParsedGrpcStatus::Success,
359    }
360
361    classify_grpc_metadata_test! {
362        name: two_success_codes_second_matches,
363        status: "3",
364        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
365        expected: ParsedGrpcStatus::Success,
366    }
367
368    classify_grpc_metadata_test! {
369        name: two_success_codes_none_matches,
370        status: "16",
371        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
372        expected: ParsedGrpcStatus::NonSuccess(NonZeroI32::new(16).unwrap()),
373    }
374}
rama_http/layer/classify/grpc_errors_as_failures.rs

rama_http/layer/classify/
grpc_errors_as_failures.rs
