Skip to main content

rama_http/layer/classify/
grpc_errors_as_failures.rs

1#![expect(
2    clippy::allow_attributes,
3    reason = "macro-generated `#[allow]` attributes whose underlying lints fire only for some expansions"
4)]
5
6use super::{ClassifiedResponse, ClassifyEos, ClassifyResponse, SharedClassifier};
7use crate::{HeaderMap, Response};
8use bitflags::bitflags;
9use percent_encoding::percent_decode;
10use std::{fmt, num::NonZeroI32};
11
12/// gRPC status codes. Used in [`GrpcErrorsAsFailures`].
13///
14/// These variants match the [gRPC status codes].
15///
16/// [gRPC status codes]: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
18#[repr(i32)]
19#[non_exhaustive]
20pub enum GrpcCode {
21    /// The operation completed successfully.
22    Ok = 0,
23    /// The operation was cancelled.
24    Cancelled = 1,
25    /// Unknown error.
26    Unknown = 2,
27    /// Client specified an invalid argument.
28    InvalidArgument = 3,
29    /// Deadline expired before operation could complete.
30    DeadlineExceeded = 4,
31    /// Some requested entity was not found.
32    NotFound = 5,
33    /// Some entity that we attempted to create already exists.
34    AlreadyExists = 6,
35    /// The caller does not have permission to execute the specified operation.
36    PermissionDenied = 7,
37    /// Some resource has been exhausted.
38    ResourceExhausted = 8,
39    /// The system is not in a state required for the operation's execution.
40    FailedPrecondition = 9,
41    /// The operation was aborted.
42    Aborted = 10,
43    /// Operation was attempted past the valid range.
44    OutOfRange = 11,
45    /// Operation is not implemented or not supported.
46    Unimplemented = 12,
47    /// Internal error.
48    Internal = 13,
49    /// The service is currently unavailable.
50    Unavailable = 14,
51    /// Unrecoverable data loss or corruption.
52    DataLoss = 15,
53    /// The request does not have valid authentication credentials
54    Unauthenticated = 16,
55}
56
57impl GrpcCode {
58    pub(crate) const fn into_bitmask(self) -> GrpcCodeBitmask {
59        match self {
60            Self::Ok => GrpcCodeBitmask::OK,
61            Self::Cancelled => GrpcCodeBitmask::CANCELLED,
62            Self::Unknown => GrpcCodeBitmask::UNKNOWN,
63            Self::InvalidArgument => GrpcCodeBitmask::INVALID_ARGUMENT,
64            Self::DeadlineExceeded => GrpcCodeBitmask::DEADLINE_EXCEEDED,
65            Self::NotFound => GrpcCodeBitmask::NOT_FOUND,
66            Self::AlreadyExists => GrpcCodeBitmask::ALREADY_EXISTS,
67            Self::PermissionDenied => GrpcCodeBitmask::PERMISSION_DENIED,
68            Self::ResourceExhausted => GrpcCodeBitmask::RESOURCE_EXHAUSTED,
69            Self::FailedPrecondition => GrpcCodeBitmask::FAILED_PRECONDITION,
70            Self::Aborted => GrpcCodeBitmask::ABORTED,
71            Self::OutOfRange => GrpcCodeBitmask::OUT_OF_RANGE,
72            Self::Unimplemented => GrpcCodeBitmask::UNIMPLEMENTED,
73            Self::Internal => GrpcCodeBitmask::INTERNAL,
74            Self::Unavailable => GrpcCodeBitmask::UNAVAILABLE,
75            Self::DataLoss => GrpcCodeBitmask::DATA_LOSS,
76            Self::Unauthenticated => GrpcCodeBitmask::UNAUTHENTICATED,
77        }
78    }
79
80    fn from_i32(code: i32) -> Option<Self> {
81        match code {
82            0 => Some(Self::Ok),
83            1 => Some(Self::Cancelled),
84            2 => Some(Self::Unknown),
85            3 => Some(Self::InvalidArgument),
86            4 => Some(Self::DeadlineExceeded),
87            5 => Some(Self::NotFound),
88            6 => Some(Self::AlreadyExists),
89            7 => Some(Self::PermissionDenied),
90            8 => Some(Self::ResourceExhausted),
91            9 => Some(Self::FailedPrecondition),
92            10 => Some(Self::Aborted),
93            11 => Some(Self::OutOfRange),
94            12 => Some(Self::Unimplemented),
95            13 => Some(Self::Internal),
96            14 => Some(Self::Unavailable),
97            15 => Some(Self::DataLoss),
98            16 => Some(Self::Unauthenticated),
99            _ => None,
100        }
101    }
102}
103
104/// Converts an `i32` gRPC status code into a [`GrpcCode`].
105///
106/// Unrecognized codes (outside 0-16) map to [`GrpcCode::Unknown`].
107impl From<i32> for GrpcCode {
108    fn from(value: i32) -> Self {
109        if value == 2 {
110            return Self::Unknown;
111        }
112
113        match value {
114            0 => Self::Ok,
115            1 => Self::Cancelled,
116            3 => Self::InvalidArgument,
117            4 => Self::DeadlineExceeded,
118            5 => Self::NotFound,
119            6 => Self::AlreadyExists,
120            7 => Self::PermissionDenied,
121            8 => Self::ResourceExhausted,
122            9 => Self::FailedPrecondition,
123            10 => Self::Aborted,
124            11 => Self::OutOfRange,
125            12 => Self::Unimplemented,
126            13 => Self::Internal,
127            14 => Self::Unavailable,
128            15 => Self::DataLoss,
129            16 => Self::Unauthenticated,
130            _ => Self::Unknown,
131        }
132    }
133}
134
135impl From<NonZeroI32> for GrpcCode {
136    fn from(value: NonZeroI32) -> Self {
137        Self::from(value.get())
138    }
139}
140
141bitflags! {
142    #[derive(Debug, Clone, Copy)]
143    pub(crate) struct GrpcCodeBitmask: u32 {
144        const OK                  = 0b00000000000000001;
145        const CANCELLED           = 0b00000000000000010;
146        const UNKNOWN             = 0b00000000000000100;
147        const INVALID_ARGUMENT    = 0b00000000000001000;
148        const DEADLINE_EXCEEDED   = 0b00000000000010000;
149        const NOT_FOUND           = 0b00000000000100000;
150        const ALREADY_EXISTS      = 0b00000000001000000;
151        const PERMISSION_DENIED   = 0b00000000010000000;
152        const RESOURCE_EXHAUSTED  = 0b00000000100000000;
153        const FAILED_PRECONDITION = 0b00000001000000000;
154        const ABORTED             = 0b00000010000000000;
155        const OUT_OF_RANGE        = 0b00000100000000000;
156        const UNIMPLEMENTED       = 0b00001000000000000;
157        const INTERNAL            = 0b00010000000000000;
158        const UNAVAILABLE         = 0b00100000000000000;
159        const DATA_LOSS           = 0b01000000000000000;
160        const UNAUTHENTICATED     = 0b10000000000000000;
161    }
162}
163
164impl From<GrpcCode> for GrpcCodeBitmask {
165    fn from(code: GrpcCode) -> Self {
166        match code {
167            GrpcCode::Ok => Self::OK,
168            GrpcCode::Cancelled => Self::CANCELLED,
169            GrpcCode::Unknown => Self::UNKNOWN,
170            GrpcCode::InvalidArgument => Self::INVALID_ARGUMENT,
171            GrpcCode::DeadlineExceeded => Self::DEADLINE_EXCEEDED,
172            GrpcCode::NotFound => Self::NOT_FOUND,
173            GrpcCode::AlreadyExists => Self::ALREADY_EXISTS,
174            GrpcCode::PermissionDenied => Self::PERMISSION_DENIED,
175            GrpcCode::ResourceExhausted => Self::RESOURCE_EXHAUSTED,
176            GrpcCode::FailedPrecondition => Self::FAILED_PRECONDITION,
177            GrpcCode::Aborted => Self::ABORTED,
178            GrpcCode::OutOfRange => Self::OUT_OF_RANGE,
179            GrpcCode::Unimplemented => Self::UNIMPLEMENTED,
180            GrpcCode::Internal => Self::INTERNAL,
181            GrpcCode::Unavailable => Self::UNAVAILABLE,
182            GrpcCode::DataLoss => Self::DATA_LOSS,
183            GrpcCode::Unauthenticated => Self::UNAUTHENTICATED,
184        }
185    }
186}
187
188/// Response classifier for gRPC responses.
189///
190/// gRPC doesn't use normal HTTP statuses for indicating success or failure but instead a special
191/// header that might appear in a trailer.
192///
193/// Responses are considered successful if
194///
195/// - `grpc-status` header value contains a success value.
196/// - `grpc-status` header is missing.
197/// - `grpc-status` header value isn't a valid `String`.
198/// - `grpc-status` header value can't parsed into an `i32`.
199///
200/// All others are considered failures.
201#[derive(Debug, Clone)]
202pub struct GrpcErrorsAsFailures {
203    success_codes: GrpcCodeBitmask,
204}
205
206impl Default for GrpcErrorsAsFailures {
207    fn default() -> Self {
208        Self::new()
209    }
210}
211
212impl GrpcErrorsAsFailures {
213    /// Create a new [`GrpcErrorsAsFailures`].
214    #[must_use]
215    pub const fn new() -> Self {
216        Self {
217            success_codes: GrpcCodeBitmask::OK,
218        }
219    }
220
221    rama_utils::macros::generate_set_and_with! {
222        /// Change which gRPC codes are considered success.
223        ///
224        /// Defaults to only considering `Ok` as success.
225        ///
226        /// `Ok` will always be considered a success.
227        pub fn success(mut self, code: GrpcCode) -> Self {
228            self.success_codes |= code.into_bitmask();
229            self
230        }
231    }
232
233    /// Returns a [`MakeClassifier`](super::MakeClassifier) that produces `GrpcErrorsAsFailures`.
234    ///
235    /// This is a convenience function that simply calls `SharedClassifier::new`.
236    #[must_use]
237    pub fn make_classifier() -> SharedClassifier<Self> {
238        SharedClassifier::new(Self::new())
239    }
240}
241
242impl ClassifyResponse for GrpcErrorsAsFailures {
243    type FailureClass = GrpcFailureClass;
244    type ClassifyEos = GrpcEosErrorsAsFailures;
245
246    fn classify_response<B>(
247        self,
248        res: &Response<B>,
249    ) -> ClassifiedResponse<Self::FailureClass, Self::ClassifyEos> {
250        match classify_grpc_metadata(res.headers(), self.success_codes) {
251            ParsedGrpcStatus::Success | ParsedGrpcStatus::HeaderNotGrpcCode => {
252                ClassifiedResponse::Ready(Ok(()))
253            }
254            ParsedGrpcStatus::NonSuccess(status) => {
255                ClassifiedResponse::Ready(Err(GrpcFailureClass::Status(status)))
256            }
257            ParsedGrpcStatus::GrpcStatusHeaderMissing => {
258                ClassifiedResponse::RequiresEos(GrpcEosErrorsAsFailures {
259                    success_codes: self.success_codes,
260                })
261            }
262        }
263    }
264
265    fn classify_error<E>(self, error: &E) -> Self::FailureClass
266    where
267        E: fmt::Display,
268    {
269        GrpcFailureClass::Error(error.to_string())
270    }
271}
272
273/// The [`ClassifyEos`] for [`GrpcErrorsAsFailures`].
274#[derive(Debug, Clone)]
275pub struct GrpcEosErrorsAsFailures {
276    success_codes: GrpcCodeBitmask,
277}
278
279impl ClassifyEos for GrpcEosErrorsAsFailures {
280    type FailureClass = GrpcFailureClass;
281
282    fn classify_eos(self, trailers: Option<&HeaderMap>) -> Result<(), Self::FailureClass> {
283        if let Some(trailers) = trailers {
284            match classify_grpc_metadata(trailers, self.success_codes) {
285                ParsedGrpcStatus::Success
286                | ParsedGrpcStatus::GrpcStatusHeaderMissing
287                | ParsedGrpcStatus::HeaderNotGrpcCode => Ok(()),
288                ParsedGrpcStatus::NonSuccess(status) => Err(GrpcFailureClass::Status(status)),
289            }
290        } else {
291            Ok(())
292        }
293    }
294
295    fn classify_error<E>(self, error: &E) -> Self::FailureClass
296    where
297        E: fmt::Display,
298    {
299        GrpcFailureClass::Error(error.to_string())
300    }
301}
302
303impl Default for GrpcEosErrorsAsFailures {
304    fn default() -> Self {
305        Self {
306            success_codes: GrpcCodeBitmask::OK,
307        }
308    }
309}
310
311/// The failure class for [`GrpcErrorsAsFailures`].
312#[derive(Debug)]
313#[non_exhaustive]
314pub enum GrpcFailureClass {
315    /// A gRPC response was classified as a failure with the corresponding status.
316    Status(GrpcStatus),
317    /// A gRPC response was classified as an error with the corresponding error description.
318    Error(String),
319}
320
321impl fmt::Display for GrpcFailureClass {
322    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
323        match self {
324            Self::Status(status) => write!(f, "Status: {status}"),
325            Self::Error(error) => write!(f, "Error: {error}"),
326        }
327    }
328}
329
330impl std::error::Error for GrpcFailureClass {}
331
332#[allow(clippy::match_result_ok)]
333pub(crate) fn classify_grpc_metadata(
334    headers: &HeaderMap,
335    success_codes: GrpcCodeBitmask,
336) -> ParsedGrpcStatus {
337    macro_rules! or_else {
338        ($expr:expr, $other:ident) => {
339            if let Some(value) = $expr {
340                value
341            } else {
342                return ParsedGrpcStatus::$other;
343            }
344        };
345    }
346
347    let code_header = or_else!(headers.get("grpc-status"), GrpcStatusHeaderMissing);
348    let code_value: i32 = or_else!(
349        code_header.to_str().ok().and_then(|s| s.parse().ok()),
350        HeaderNotGrpcCode
351    );
352    let grpc_code = GrpcCode::from_i32(code_value);
353
354    if let Some(code) = grpc_code
355        && success_codes.contains(GrpcCodeBitmask::from(code))
356    {
357        return ParsedGrpcStatus::Success;
358    }
359
360    let message = headers.get("grpc-message").map(|header| {
361        percent_decode(header.as_bytes())
362            .decode_utf8_lossy()
363            .into_owned()
364    });
365
366    ParsedGrpcStatus::NonSuccess(GrpcStatus {
367        code: grpc_code,
368        code_raw: code_value,
369        message,
370    })
371}
372
373/// A gRPC status extracted from response headers/trailers.
374#[derive(Debug, PartialEq, Eq, Hash)]
375pub struct GrpcStatus {
376    code: Option<GrpcCode>,
377    code_raw: i32,
378    message: Option<String>,
379}
380
381impl GrpcStatus {
382    /// Returns the status code as a [`GrpcCode`], or `None` if the code is not recognized.
383    pub fn code(&self) -> Option<GrpcCode> {
384        self.code
385    }
386
387    /// Returns the raw integer status code.
388    pub fn code_raw(&self) -> i32 {
389        self.code_raw
390    }
391
392    /// Returns the percent-decoded gRPC error message, if present.
393    pub fn message(&self) -> Option<&str> {
394        self.message.as_deref()
395    }
396}
397
398impl fmt::Display for GrpcStatus {
399    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
400        match self.code {
401            Some(code) => write!(f, "{code:?}")?,
402            None => write!(f, "Code({})", self.code_raw)?,
403        }
404        if let Some(message) = self.message.as_ref() {
405            write!(f, ": {message}")?;
406        }
407        Ok(())
408    }
409}
410
411#[derive(Debug, PartialEq, Eq, Hash)]
412pub(crate) enum ParsedGrpcStatus {
413    Success,
414    NonSuccess(GrpcStatus),
415    GrpcStatusHeaderMissing,
416    // this is treated as `Success` but kept separate for clarity
417    HeaderNotGrpcCode,
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    macro_rules! classify_grpc_metadata_test {
425        (
426            name: $name:ident,
427            status: $status:expr,
428            success_flags: $success_flags:expr,
429            expected: $expected:expr,
430        ) => {
431            classify_grpc_metadata_test!(
432                name: $name,
433                status: $status,
434                message: "",
435                success_flags: $success_flags,
436                expected: $expected,
437            );
438        };
439        (
440            name: $name:ident,
441            status: $status:expr,
442            message: $message:expr,
443            success_flags: $success_flags:expr,
444            expected: $expected:expr,
445        ) => {
446            #[test]
447            fn $name() {
448                let mut headers = HeaderMap::new();
449                headers.insert("grpc-status", $status.parse().unwrap());
450                if !$message.is_empty() {
451                    headers.insert("grpc-message", $message.parse().unwrap());
452                }
453                let status = classify_grpc_metadata(&headers, $success_flags);
454                assert_eq!(status, $expected);
455            }
456        };
457    }
458
459    classify_grpc_metadata_test! {
460        name: basic_ok,
461        status: "0",
462        success_flags: GrpcCodeBitmask::OK,
463        expected: ParsedGrpcStatus::Success,
464    }
465
466    classify_grpc_metadata_test! {
467        name: basic_error,
468        status: "1",
469        success_flags: GrpcCodeBitmask::OK,
470        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
471            code: Some(GrpcCode::Cancelled),
472            code_raw: 1,
473            message: None,
474        }),
475    }
476
477    classify_grpc_metadata_test! {
478        name: two_success_codes_first_matches,
479        status: "0",
480        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
481        expected: ParsedGrpcStatus::Success,
482    }
483
484    classify_grpc_metadata_test! {
485        name: two_success_codes_second_matches,
486        status: "3",
487        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
488        expected: ParsedGrpcStatus::Success,
489    }
490
491    classify_grpc_metadata_test! {
492        name: two_success_codes_none_matches,
493        status: "16",
494        message: "mock message",
495        success_flags: GrpcCodeBitmask::OK | GrpcCodeBitmask::INVALID_ARGUMENT,
496        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
497            code: Some(GrpcCode::Unauthenticated),
498            code_raw: 16,
499            message: Some("mock message".to_owned()),
500        }),
501    }
502
503    classify_grpc_metadata_test! {
504        name: percent_encoded_message,
505        status: "2",
506        message: "hello%20world",
507        success_flags: GrpcCodeBitmask::OK,
508        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
509            code: Some(GrpcCode::Unknown),
510            code_raw: 2,
511            message: Some("hello world".to_owned()),
512        }),
513    }
514
515    classify_grpc_metadata_test! {
516        name: invalid_percent_encoding,
517        status: "13",
518        message: "bad%2Gencode",
519        success_flags: GrpcCodeBitmask::OK,
520        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
521            code: Some(GrpcCode::Internal),
522            code_raw: 13,
523            message: Some("bad%2Gencode".to_owned()),
524        }),
525    }
526
527    classify_grpc_metadata_test! {
528        name: empty_grpc_message,
529        status: "5",
530        message: "",
531        success_flags: GrpcCodeBitmask::OK,
532        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
533            code: Some(GrpcCode::NotFound),
534            code_raw: 5,
535            message: None,
536        }),
537    }
538
539    classify_grpc_metadata_test! {
540        name: unknown_status_code_above_16,
541        status: "99",
542        message: "custom error",
543        success_flags: GrpcCodeBitmask::OK,
544        expected: ParsedGrpcStatus::NonSuccess(GrpcStatus{
545            code: None,
546            code_raw: 99,
547            message: Some("custom error".to_owned()),
548        }),
549    }
550
551    #[test]
552    fn invalid_utf8_after_percent_decode() {
553        let mut headers = HeaderMap::new();
554        headers.insert("grpc-status", "2".parse().unwrap());
555        // %80 is an invalid UTF-8 start byte; lossy decode replaces it with U+FFFD
556        headers.insert("grpc-message", "bad%80byte".parse().unwrap());
557        let status = classify_grpc_metadata(&headers, GrpcCodeBitmask::OK);
558        assert_eq!(
559            status,
560            ParsedGrpcStatus::NonSuccess(GrpcStatus {
561                code: Some(GrpcCode::Unknown),
562                code_raw: 2,
563                message: Some("bad\u{FFFD}byte".to_owned()),
564            })
565        );
566    }
567
568    #[test]
569    fn valid_utf8_percent_encoded() {
570        let mut headers = HeaderMap::new();
571        headers.insert("grpc-status", "3".parse().unwrap());
572        // %C3%A9 is the percent-encoded form of 'é' (U+00E9) in UTF-8
573        headers.insert("grpc-message", "caf%C3%A9".parse().unwrap());
574        let status = classify_grpc_metadata(&headers, GrpcCodeBitmask::OK);
575        assert_eq!(
576            status,
577            ParsedGrpcStatus::NonSuccess(GrpcStatus {
578                code: Some(GrpcCode::InvalidArgument),
579                code_raw: 3,
580                message: Some("café".to_owned()),
581            })
582        );
583    }
584
585    #[test]
586    fn grpc_ok_classified_as_success() {
587        let res = Response::builder()
588            .header("grpc-status", "0")
589            .body(())
590            .unwrap();
591
592        let classifier = GrpcErrorsAsFailures::new();
593        let result = classifier.classify_response(&res);
594        assert!(matches!(result, ClassifiedResponse::Ready(Ok(()))));
595    }
596
597    #[test]
598    fn grpc_code_from_i32_known_codes() {
599        assert!(matches!(GrpcCode::from(0), GrpcCode::Ok));
600        assert!(matches!(GrpcCode::from(1), GrpcCode::Cancelled));
601        assert!(matches!(GrpcCode::from(4), GrpcCode::DeadlineExceeded));
602        assert!(matches!(GrpcCode::from(13), GrpcCode::Internal));
603        assert!(matches!(GrpcCode::from(16), GrpcCode::Unauthenticated));
604    }
605
606    #[test]
607    fn grpc_code_from_i32_unknown_codes() {
608        assert!(matches!(GrpcCode::from(17), GrpcCode::Unknown));
609        assert!(matches!(GrpcCode::from(-1), GrpcCode::Unknown));
610        assert!(matches!(GrpcCode::from(9999), GrpcCode::Unknown));
611    }
612
613    #[test]
614    fn grpc_code_from_non_zero_i32() {
615        let code = NonZeroI32::new(7).unwrap();
616        assert!(matches!(GrpcCode::from(code), GrpcCode::PermissionDenied));
617
618        let code = NonZeroI32::new(99).unwrap();
619        assert!(matches!(GrpcCode::from(code), GrpcCode::Unknown));
620    }
621}