refactor(api): Simplify Protobuf codec

Author: alixinne

src/servers/proto.rs

@@ -3,74 +3,45 @@
 use std::net::SocketAddr;
 
 use futures::prelude::*;
-use prost::Message;
 use thiserror::Error;
 use tokio::net::TcpStream;
+use tokio_util::codec::Framed;
 
 use crate::{
     api::proto::{self, message, ProtoApiError},
-    global::{Global, InputMessage, InputSourceHandle, InputSourceName},
+    global::{Global, InputSourceName},
 };
 
+mod codec;
+use codec::*;
+
 #[derive(Debug, Error)]
 pub enum ProtoServerError {
     #[error("i/o error: {0}")]
     Io(#[from] futures_io::Error),
-    #[error("decode error: {}", 0)]
-    DecodeError(#[from] prost::DecodeError),
+    #[error("decode error: {0}")]
+    Codec(#[from] ProtoCodecError),
     #[error(transparent)]
     Api(#[from] ProtoApiError),
 }
 
-fn encode_response(buf: &mut bytes::BytesMut, msg: impl prost::Message) -> bytes::Bytes {
-    // Clear the buffer to start fresh
-    buf.clear();
-
-    // Reserve enough space for the response
-    let len = msg.encoded_len();
-    if buf.capacity() < len {
-        buf.reserve(len * 2);
-    }
-
-    // Encode the message
-    msg.encode(buf).unwrap();
-    buf.split().freeze()
-}
-
-fn success_response(peer_addr: SocketAddr, buf: &mut bytes::BytesMut) -> bytes::Bytes {
+fn success_response(peer_addr: SocketAddr) -> message::HyperionReply {
     let mut reply = message::HyperionReply::default();
     reply.r#type = message::hyperion_reply::Type::Reply.into();
     reply.success = Some(true);
 
     trace!("({}) sending success: {:?}", peer_addr, reply);
-    encode_response(buf, reply)
+    reply
 }
 
-fn error_response(
-    peer_addr: SocketAddr,
-    buf: &mut bytes::BytesMut,
-    error: impl std::fmt::Display,
-) -> bytes::Bytes {
+fn error_response(peer_addr: SocketAddr, error: impl std::fmt::Display) -> message::HyperionReply {
     let mut reply = message::HyperionReply::default();
     reply.r#type = message::hyperion_reply::Type::Reply.into();
     reply.success = Some(false);
     reply.error = Some(error.to_string());
 
     trace!("({}) sending error: {:?}", peer_addr, reply);
-    encode_response(buf, reply)
-}
-
-fn handle_request(
-    peer_addr: SocketAddr,
-    request_bytes: bytes::BytesMut,
-    source: &InputSourceHandle<InputMessage>,
-) -> Result<(), ProtoServerError> {
-    let request_bytes = request_bytes.freeze();
-    let request = message::HyperionRequest::decode(request_bytes.clone())?;
-
-    trace!("({}) got request: {:?}", peer_addr, request);
-
-    Ok(proto::handle_request(request, source)?)
+    reply
 }
 
 pub async fn handle_client(
@@ -79,35 +50,31 @@ pub async fn handle_client(
 ) -> Result<(), ProtoServerError> {
     debug!("accepted new connection from {}", peer_addr);
 
-    let framed = tokio_util::codec::LengthDelimitedCodec::builder()
-        .length_field_length(4)
-        .new_framed(socket);
-    let (mut writer, mut reader) = framed.split();
+    let (mut writer, mut reader) = Framed::new(socket, ProtoCodec::new()).split();
 
     // unwrap: cannot fail because the priority is None
     let source = global
         .register_input_source(InputSourceName::Protobuf { peer_addr }, None)
         .await
         .unwrap();
 
-    // buffer for building responses
-    let mut reply_buf = bytes::BytesMut::with_capacity(128);
-
-    while let Some(request_bytes) = reader.next().await {
-        let request_bytes = match request_bytes {
+    while let Some(request) = reader.next().await {
+        let request = match request {
             Ok(rb) => rb,
             Err(error) => {
                 error!("({}) error reading frame: {}", peer_addr, error);
                 continue;
             }
         };
 
-        let reply = match handle_request(peer_addr, request_bytes, &source) {
-            Ok(()) => success_response(peer_addr, &mut reply_buf),
+        trace!("({}) got request: {:?}", peer_addr, request);
+
+        let reply = match proto::handle_request(request, &source) {
+            Ok(()) => success_response(peer_addr),
             Err(error) => {
                 error!("({}) error processing request: {}", peer_addr, error);
 
-                error_response(peer_addr, &mut reply_buf, error)
+                error_response(peer_addr, error)
             }
         };
 

src/servers/proto/codec.rs

@@ -0,0 +1,71 @@
+use bytes::BytesMut;
+use prost::Message;
+use thiserror::Error;
+use tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec};
+
+use crate::api::proto::message;
+
+#[derive(Debug, Error)]
+pub enum ProtoCodecError {
+    #[error("i/o error: {0}")]
+    Io(#[from] futures_io::Error),
+    #[error(transparent)]
+    LengthDelimited(#[from] tokio_util::codec::LengthDelimitedCodecError),
+    #[error(transparent)]
+    Decode(#[from] prost::DecodeError),
+    #[error(transparent)]
+    Encode(#[from] prost::EncodeError),
+}
+
+/// JSON tokio codec
+pub struct ProtoCodec {
+    /// Line parsing codec
+    inner: LengthDelimitedCodec,
+    /// Buffer for encoding messages
+    buf: BytesMut,
+}
+
+impl ProtoCodec {
+    /// Create a new ProtoCodec
+    pub fn new() -> Self {
+        Self {
+            inner: LengthDelimitedCodec::builder()
+                .length_field_length(4)
+                .new_codec(),
+            buf: BytesMut::new(),
+        }
+    }
+}
+
+impl Decoder for ProtoCodec {
+    type Item = message::HyperionRequest;
+    type Error = ProtoCodecError;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        match self.inner.decode(src) {
+            Ok(inner_result) => Ok(match inner_result {
+                Some(ref data) => Some(message::HyperionRequest::decode(data.clone().freeze())?),
+                None => None,
+            }),
+            Err(error) => Err(error.into()),
+        }
+    }
+}
+
+impl Encoder<message::HyperionReply> for ProtoCodec {
+    type Error = ProtoCodecError;
+
+    fn encode(
+        &mut self,
+        item: message::HyperionReply,
+        dst: &mut BytesMut,
+    ) -> Result<(), Self::Error> {
+        self.buf.clear();
+        self.buf.reserve(item.encoded_len());
+
+        match item.encode(&mut self.buf) {
+            Ok(_) => Ok(self.inner.encode(self.buf.clone().freeze(), dst)?),
+            Err(encode_error) => Err(encode_error.into()),
+        }
+    }
+}