Use Rust types to force IntoTwirpResponse to return a valid body.

jorendorff · jorendorff · commit 319d675104b4 · 2025-01-13T14:24:53.000-06:00
It's still possible to return a response with an invalid status code or headers, but less likely to
happen by accident.
diff --git a/crates/twirp/src/error.rs b/crates/twirp/src/error.rs
@@ -17,11 +17,11 @@ pub trait IntoTwirpResponse {
     /// ```
     /// use axum::body::Body;
     /// use http::Response;
-    /// use twirp::IntoTwirpResponse;
+    /// use twirp::{TwirpErrorResponse, IntoTwirpResponse};
     /// # struct MyError { message: String }
     ///
     /// impl IntoTwirpResponse for MyError {
-    ///     fn into_twirp_response(self) -> Response<Body> {
+    ///     fn into_twirp_response(self) -> Response<TwirpErrorResponse> {
     ///         // Use TwirpErrorResponse to generate a valid starting point
     ///         let mut response = twirp::invalid_argument(&self.message)
     ///             .into_twirp_response();
@@ -35,7 +35,7 @@ pub trait IntoTwirpResponse {
     ///
     /// The `Response` that `TwirpErrorResponse` generates can be used as a starting point,
     /// adding headers and extensions to it.
-    fn into_twirp_response(self) -> Response<Body>;
+    fn into_twirp_response(self) -> Response<TwirpErrorResponse>;
 }
 
 /// Alias for a generic error
@@ -182,26 +182,30 @@ impl TwirpErrorResponse {
     pub fn insert_meta(&mut self, key: String, value: String) -> Option<String> {
         self.meta.insert(key, value)
     }
-}
 
-impl IntoTwirpResponse for TwirpErrorResponse {
-    fn into_twirp_response(self) -> Response<Body> {
-        IntoResponse::into_response(self)
+    pub fn into_axum_body(self) -> Body {
+        let json =
+            serde_json::to_string(&self).expect("JSON serialization of an error should not fail");
+        Body::new(json)
     }
 }
 
-impl IntoResponse for TwirpErrorResponse {
-    fn into_response(self) -> Response<Body> {
+impl IntoTwirpResponse for TwirpErrorResponse {
+    fn into_twirp_response(self) -> Response<TwirpErrorResponse> {
         let mut headers = HeaderMap::new();
         headers.insert(
             header::CONTENT_TYPE,
             HeaderValue::from_static("application/json"),
         );
 
-        let json =
-            serde_json::to_string(&self).expect("JSON serialization of an error should not fail");
+        let code = self.code.http_status_code();
+        (code, headers).into_response().map(|_| self)
+    }
+}
 
-        (self.code.http_status_code(), headers, json).into_response()
+impl IntoResponse for TwirpErrorResponse {
+    fn into_response(self) -> Response<Body> {
+        self.into_twirp_response().map(|err| err.into_axum_body())
     }
 }
 
diff --git a/crates/twirp/src/server.rs b/crates/twirp/src/server.rs
@@ -70,7 +70,7 @@ where
             //     .insert(RequestError(err));
             let mut twirp_err = error::malformed("bad request");
             twirp_err.insert_meta("error".to_string(), err.to_string());
-            return twirp_err.into_twirp_response();
+            return twirp_err.into_response();
         }
     };
 
@@ -85,7 +85,7 @@ where
             // TODO: Capture original error in the response extensions.
             let mut twirp_err = error::unknown("error serializing response");
             twirp_err.insert_meta("error".to_string(), err.to_string());
-            return twirp_err.into_twirp_response();
+            return twirp_err.into_response();
         }
     };
     timings.set_response_written();
@@ -135,7 +135,7 @@ where
                     .body(Body::from(data))?
             }
         },
-        Err(err) => err.into_twirp_response(),
+        Err(err) => err.into_response(),
     };
     Ok(res)
 }
diff --git a/example/src/bin/advanced-server.rs b/example/src/bin/advanced-server.rs
@@ -8,7 +8,7 @@ use twirp::axum::body::Body;
 use twirp::axum::http;
 use twirp::axum::middleware::{self, Next};
 use twirp::axum::routing::get;
-use twirp::{invalid_argument, Context, IntoTwirpResponse, Router};
+use twirp::{invalid_argument, Context, IntoTwirpResponse, Router, TwirpErrorResponse};
 
 pub mod service {
     pub mod haberdash {
@@ -56,10 +56,7 @@ enum HatError {
 }
 
 impl IntoTwirpResponse for HatError {
-    fn into_twirp_response(self) -> http::Response<Body> {
-        // Note: When converting a HatError to a response, since we want the server to be a twirp
-        // server, it's important to generate a response that follows the twirp standard. We do
-        // that here by using TwirpErrorResponse.
+    fn into_twirp_response(self) -> http::Response<TwirpErrorResponse> {
         match self {
             HatError::InvalidSize => invalid_argument("inches").into_twirp_response(),
         }

Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,7 @@ where`
`70`	`70`	`// .insert(RequestError(err));`
`71`	`71`	`let mut twirp_err = error::malformed("bad request");`
`72`	`72`	`twirp_err.insert_meta("error".to_string(), err.to_string());`
`73`		`- return twirp_err.into_twirp_response();`
	`73`	`+ return twirp_err.into_response();`
`74`	`74`	`}`
`75`	`75`	`};`
`76`	`76`
`@@ -85,7 +85,7 @@ where`
`85`	`85`	`// TODO: Capture original error in the response extensions.`
`86`	`86`	`let mut twirp_err = error::unknown("error serializing response");`
`87`	`87`	`twirp_err.insert_meta("error".to_string(), err.to_string());`
`88`		`- return twirp_err.into_twirp_response();`
	`88`	`+ return twirp_err.into_response();`
`89`	`89`	`}`
`90`	`90`	`};`
`91`	`91`	`timings.set_response_written();`
`@@ -135,7 +135,7 @@ where`
`135`	`135`	`.body(Body::from(data))?`
`136`	`136`	`}`
`137`	`137`	`},`
`138`		`- Err(err) => err.into_twirp_response(),`
	`138`	`+ Err(err) => err.into_response(),`
`139`	`139`	`};`
`140`	`140`	`Ok(res)`
`141`	`141`	`}`