1 files changed, 14 insertions, 28 deletions

diff --git a/src/server.rs b/src/server.rs
index 7bf17df..eb70057 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -29,6 +29,7 @@ use kuska_handshake::async_std::{handshake_server, BoxStream};
 use crate::error::*;
 use crate::netapp::*;
 use crate::proto::*;
+use crate::proto2::*;
 use crate::util::*;
 
 // The client and server connection structs (client.rs and server.rs)
@@ -116,22 +117,8 @@ impl ServerConn {
 	}
 
 	async fn recv_handler_aux(self: &Arc<Self>, bytes: &[u8]) -> Result<Vec<u8>, Error> {
-		if bytes.len() < 2 {
-			return Err(Error::Message("Invalid protocol message".into()));
-		}
-
-		// byte 0 is the request priority, we don't care here
-		let path_length = bytes[1] as usize;
-		if bytes.len() < 2 + path_length {
-			return Err(Error::Message("Invalid protocol message".into()));
-		}
-
-		let path = &bytes[2..2 + path_length];
-		let path = String::from_utf8(path.to_vec())?;
-
-		let telemetry_id_len = bytes[2 + path_length] as usize;
-
-		let data = &bytes[3 + path_length + telemetry_id_len..];
+		let msg = QueryMessage::decode(bytes)?;
+		let path = String::from_utf8(msg.path.to_vec())?;
 
 		let handler_opt = {
 			let endpoints = self.netapp.endpoints.read().unwrap();
@@ -143,10 +130,9 @@ impl ServerConn {
 				if #[cfg(feature = "telemetry")] {
 					let tracer = opentelemetry::global::tracer("netapp");
 
-					let mut span = if telemetry_id_len > 0 {
-						let by = bytes[3+path_length..3+path_length+telemetry_id_len].to_vec();
+					let mut span = if let Some(telemetry_id) = msg.telemetry_id {
 						let propagator = BinaryPropagator::new();
-						let context = propagator.from_bytes(by);
+						let context = propagator.from_bytes(telemetry_id);
 						let context = Context::new().with_remote_span_context(context);
 						tracer.span_builder(format!(">> RPC {}", path))
 							.with_kind(SpanKind::Server)
@@ -161,13 +147,13 @@ impl ServerConn {
 							.start(&tracer)
 					};
 					span.set_attribute(KeyValue::new("path", path.to_string()));
-					span.set_attribute(KeyValue::new("len_query", data.len() as i64));
+					span.set_attribute(KeyValue::new("len_query", msg.body.len() as i64));
 
-					handler.handle(data, self.peer_id)
+					handler.handle(msg.body, self.peer_id)
 						.with_context(Context::current_with_span(span))
 						.await
 				} else {
-					handler.handle(data, self.peer_id).await
+					handler.handle(msg.body, self.peer_id).await
 				}
 			}
 		} else {
@@ -191,16 +177,16 @@ impl RecvLoop for ServerConn {
 			let prio = if !bytes.is_empty() { bytes[0] } else { 0u8 };
 			let resp = self2.recv_handler_aux(&bytes[..]).await;
 
-			let mut resp_bytes = vec![];
-			match resp {
+			let resp_bytes = match resp {
 				Ok(rb) => {
-					resp_bytes.push(0u8);
-					resp_bytes.extend(&rb[..]);
+					let mut resp_bytes = vec![0u8];
+					resp_bytes.extend(rb);
+					resp_bytes
 				}
 				Err(e) => {
-					resp_bytes.push(e.code());
+					vec![e.code()]
 				}
-			}
+			};
 
 			trace!("ServerConn sending response to {}: ", id);