balancer: make json parsing errors include the path to the key causing the error (#1345)

Author: dyc3

Cargo.lock

@@ -26,6 +26,7 @@ rand = "0.8.5"
 reqwest = { version = "0.11.17", features = ["json", "stream", "rustls-tls"] }
 serde = { version = "1", features = ["derive", "rc"] }
 serde_json = { version = "1", features = ["raw_value"] }
+serde_path_to_error = "0.1.15"
 test-context = "0.1.4"
 tokio = { version = "1", features = ["full", "tracing"] }
 tokio-tungstenite = "0.21.0"

Cargo.toml

@@ -19,6 +19,7 @@ reqwest.workspace = true
 rand.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+serde_path_to_error.workspace = true
 test-context.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true

crates/harness/Cargo.toml

@@ -270,7 +270,8 @@ impl Monolith {
             .iter()
             .filter_map(|msg| match msg {
                 Message::Text(msg) => {
-                    let msg: MsgB2M = serde_json::from_str(msg).unwrap();
+                    let jd = &mut serde_json::Deserializer::from_str(msg);
+                    let msg: MsgB2M = serde_path_to_error::deserialize(jd).unwrap();
                     Some(msg)
                 }
                 _ => None,

crates/harness/src/monolith.rs

@@ -3,8 +3,6 @@ name = "ott-balancer"
 version = "0.10.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
@@ -20,6 +18,7 @@ rand.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+serde_path_to_error.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 tokio.workspace = true

crates/ott-balancer/Cargo.toml

@@ -5,7 +5,7 @@ use tokio::sync::broadcast::error::RecvError;
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
 use tokio_tungstenite::tungstenite::Message;
-use tracing::{debug, error, info, trace};
+use tracing::{debug, error, info, trace, warn};
 use uuid::Uuid;
 
 use crate::balancer::BalancerLink;
@@ -147,7 +147,21 @@ pub async fn client_entry<'r>(
     let mut client_link;
     match message {
         Message::Text(text) => {
-            let message: ClientMessage = serde_json::from_str(&text).unwrap();
+            let jd = &mut serde_json::Deserializer::from_str(&text);
+            let message: ClientMessage = match serde_path_to_error::deserialize(jd) {
+                Ok(msg) => msg,
+                Err(err) => {
+                    warn!("failed to deserialize client message: {:?}", err);
+                    stream
+                        .close(Some(CloseFrame {
+                            code: CloseCode::Protocol,
+                            reason: "failed to deserialize message".into(),
+                        }))
+                        .await?;
+                    return Ok(());
+                }
+            };
+
             match message {
                 ClientMessage::Auth(message) => {
                     debug!("client authenticated, handing off to balancer");

crates/ott-balancer/src/client.rs

@@ -126,7 +126,21 @@ async fn connect_and_maintain(
         let monolith_id;
         match message {
             Message::Text(text) => {
-                let message: MsgM2B = serde_json::from_str(&text).unwrap();
+                let jd = &mut serde_json::Deserializer::from_str(&text);
+                let message: MsgM2B = match serde_path_to_error::deserialize(jd) {
+                    Ok(msg) => msg,
+                    Err(err) => {
+                        warn!("failed to deserialize monolith message: {:?}", err);
+                        let _ = stream
+                            .close(Some(CloseFrame {
+                                code: CloseCode::Protocol,
+                                reason: "failed to deserialize message".into(),
+                            }))
+                            .await;
+                        return;
+                    }
+                };
+
                 match message {
                     MsgM2B::Init(init) => {
                         monolith_id = init.id;
@@ -154,7 +168,7 @@ async fn connect_and_maintain(
                     _ => {
                         let _ = stream
                             .close(Some(CloseFrame {
-                                code: CloseCode::Library(4000),
+                                code: CloseCode::Protocol,
                                 reason: "did not send init".into(),
                             }))
                             .await;

crates/ott-balancer/src/connection.rs

@@ -80,9 +80,13 @@ async fn do_harness_discovery(
                         error!("expected text message from harness, got something else");
                         continue;
                     };
-                    let Ok(msg) = serde_json::from_str(text.as_str()) else {
-                        error!("failed to deserialize message from harness");
-                        continue;
+                    let jd = &mut serde_json::Deserializer::from_str(&text);
+                    let msg = match serde_path_to_error::deserialize(jd) {
+                        Ok(msg) => msg,
+                        Err(err) => {
+                            error!("failed to deserialize message from harness: {:?}", err);
+                            continue;
+                        }
                     };
                     let mut monoliths = monoliths.lock().await;
                     *monoliths = msg;

crates/ott-balancer/src/discovery/harness.rs

@@ -13,7 +13,8 @@ impl<'de> SocketMessage {
         match self {
             SocketMessage::Message(msg) => {
                 let text = msg.to_text()?;
-                let obj = serde_json::from_str(text)?;
+                let jd = &mut serde_json::Deserializer::from_str(text);
+                let obj = serde_path_to_error::deserialize(jd)?;
                 Ok(obj)
             }
             #[allow(deprecated)]