fix(tonic): Handle error on codec creation error

Avoid call to default() (which in our case is unimplemented!())
JustRustThings · May 17, 2023 · c51c8b4 · c51c8b4
1 parent 6d1579c
commit c51c8b4
Showing 1 changed file with 32 additions and 20 deletions.
diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs
@@ -414,11 +414,14 @@ fn generate_unary<T: Method>(
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
     let codec_constructor = if method.codec_path().ends_with("ProstAesCodec") {
-        quote! { #codec_name::try_from(agent_id).unwrap_or_else(|err| {
-            log::error!("Invalid agent ID: {}", err);
-            // TODO: return a 502 error (or a 200 error with grpc-status: INVALID_ARGUMENT) instead
-            #codec_name::default()
-            })
+        quote! {
+            match #codec_name::try_from(agent_id) {
+                Ok(codec) => codec,
+                Err(err) => {
+                    log::error!("Invalid agent ID: {}", err);
+                    return Ok(tonic::Status::unauthenticated("Invalid agent-id").to_http())
+                }
+            }
         }
     } else {
         quote! { #codec_name::default() }
@@ -480,11 +483,14 @@ fn generate_server_streaming<T: Method>(
     let response_stream = quote::format_ident!("{}Stream", method.identifier());
 
     let codec_constructor = if method.codec_path().ends_with("ProstAesCodec") {
-        quote! { #codec_name::try_from(agent_id).unwrap_or_else(|err| {
-            log::error!("Invalid agent ID: {}", err);
-            // TODO: return a 502 error (or a 200 error with grpc-status: INVALID_ARGUMENT) instead
-            #codec_name::default()
-            })
+        quote! {
+            match #codec_name::try_from(agent_id) {
+                Ok(codec) => codec,
+                Err(err) => {
+                    log::error!("Invalid agent ID: {}", err);
+                    return Ok(tonic::Status::unauthenticated("Invalid agent-id").to_http())
+                }
+            }
         }
     } else {
         quote! { #codec_name::default() }
@@ -543,11 +549,14 @@ fn generate_client_streaming<T: Method>(
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
     let codec_constructor = if method.codec_path().ends_with("ProstAesCodec") {
-        quote! { #codec_name::try_from(agent_id).unwrap_or_else(|err| {
-            log::error!("Invalid agent ID: {}", err);
-            // TODO: return a 502 error (or a 200 error with grpc-status: INVALID_ARGUMENT) instead
-            #codec_name::default()
-            })
+        quote! {
+            match #codec_name::try_from(agent_id) {
+                Ok(codec) => codec,
+                Err(err) => {
+                    log::error!("Invalid agent ID: {}", err);
+                    return Ok(tonic::Status::unauthenticated("Invalid agent-id").to_http())
+                }
+            }
         }
     } else {
         quote! { #codec_name::default() }
@@ -611,11 +620,14 @@ fn generate_streaming<T: Method>(
     let response_stream = quote::format_ident!("{}Stream", method.identifier());
 
     let codec_constructor = if method.codec_path().ends_with("ProstAesCodec") {
-        quote! { #codec_name::try_from(agent_id).unwrap_or_else(|err| {
-            log::error!("Invalid agent ID: {}", err);
-            // TODO: return a 502 error (or a 200 error with grpc-status: INVALID_ARGUMENT) instead
-            #codec_name::default()
-            })
+        quote! {
+            match #codec_name::try_from(agent_id) {
+                Ok(codec) => codec,
+                Err(err) => {
+                    log::error!("Invalid agent ID: {}", err);
+                    return Ok(tonic::Status::unauthenticated("Invalid agent-id").to_http())
+                }
+            }
         }
     } else {
         quote! { #codec_name::default() }