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#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
18#[repr(i32)]
19#[non_exhaustive]
20pub enum GrpcCode {
21 Ok = 0,
23 Cancelled = 1,
25 Unknown = 2,
27 InvalidArgument = 3,
29 DeadlineExceeded = 4,
31 NotFound = 5,
33 AlreadyExists = 6,
35 PermissionDenied = 7,
37 ResourceExhausted = 8,
39 FailedPrecondition = 9,
41 Aborted = 10,
43 OutOfRange = 11,
45 Unimplemented = 12,
47 Internal = 13,
49 Unavailable = 14,
51 DataLoss = 15,
53 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
104impl 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#[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 #[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 pub fn success(mut self, code: GrpcCode) -> Self {
228 self.success_codes |= code.into_bitmask();
229 self
230 }
231 }
232
233 #[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#[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#[derive(Debug)]
313#[non_exhaustive]
314pub enum GrpcFailureClass {
315 Status(GrpcStatus),
317 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#[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 pub fn code(&self) -> Option<GrpcCode> {
384 self.code
385 }
386
387 pub fn code_raw(&self) -> i32 {
389 self.code_raw
390 }
391
392 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 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 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 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}