Add decode_request_header_from_buffer and encode_request_header_to_buffer helper methods (#131)

Co-authored-by: enzolombardi <enzo.lombardi@cleafy.com>

Author: aovestdipaperino

src/lib.rs

@@ -13,16 +13,17 @@
 //! easy serialization and deserialization via [`bytes::Bytes`].
 //!
 //! ```rust
-//! use bytes::{Bytes, BytesMut};
+//! use bytes::BytesMut;
+//! use kafka_protocol::protocol::{encode_request_header_into_buffer, decode_request_header_from_buffer};
 //! use kafka_protocol::messages::RequestHeader;
 //! use kafka_protocol::protocol::{StrBytes, Encodable, Decodable};
 //!
 //! let mut request_header = RequestHeader::default();
 //! request_header.correlation_id = 1;
 //! request_header.client_id = Some(StrBytes::from_static_str("test-client"));
 //! let mut buf = BytesMut::new();
-//! request_header.encode(&mut buf, 2);
-//! assert_eq!(request_header, RequestHeader::decode(&mut buf, 2).unwrap());
+//! encode_request_header_into_buffer(&mut buf, &request_header).unwrap();
+//! assert_eq!(request_header, decode_request_header_from_buffer(&mut buf).unwrap());
 //! ```
 //! Note that every message implementation of [`Encodable::encode`](crate::protocol::Encodable::encode)
 //! and [`Decodable::decode`](crate::protocol::Decodable::decode) requires a version to be provided
@@ -37,7 +38,7 @@
 //! type to a [`bytes::Bytes`] buffer.
 //!
 //! ```rust
-//! use kafka_protocol::protocol::{StrBytes, Encodable, HeaderVersion};
+//! use kafka_protocol::protocol::{StrBytes, Encodable, HeaderVersion, encode_request_header_into_buffer};
 //! use bytes::{BytesMut, Bytes};
 //! use kafka_protocol::messages::{RequestHeader, ApiKey, ApiVersionsRequest};
 //! # use std::error::Error;
@@ -47,7 +48,7 @@
 //! req_header.request_api_version = 3;
 //! req_header.request_api_key = ApiKey::ApiVersions as i16;
 //! req_header.client_id = Some(StrBytes::from_static_str("example"));
-//! req_header.encode(&mut buf, ApiVersionsRequest::header_version(req_header.request_api_version)).unwrap();
+//! encode_request_header_into_buffer(&mut buf, &req_header).unwrap();
 //! let mut api_versions_req = ApiVersionsRequest::default();
 //! api_versions_req.client_software_version = StrBytes::from_static_str("1.0");
 //! api_versions_req.client_software_name = StrBytes::from_static_str("example-client");
@@ -69,7 +70,7 @@
 //! A simple example for decoding an unknown message encoded in `buf`:
 //! ```rust
 //! use kafka_protocol::messages::{RequestHeader, ApiVersionsRequest, ApiKey, RequestKind};
-//! use kafka_protocol::protocol::{Encodable, Decodable, StrBytes, HeaderVersion};
+//! use kafka_protocol::protocol::{Encodable, Decodable, StrBytes, HeaderVersion, decode_request_header_from_buffer, encode_request_header_into_buffer};
 //! use bytes::{BytesMut, Buf};
 //! use std::convert::TryFrom;
 //! use kafka_protocol::protocol::buf::ByteBuf;
@@ -78,17 +79,14 @@
 //! # req_header.request_api_version = 3;
 //! # req_header.request_api_key = ApiKey::ApiVersions as i16;
 //! # req_header.client_id = Some(StrBytes::from_static_str("example"));
-//! # req_header.encode(&mut buf, ApiVersionsRequest::header_version(req_header.request_api_version)).unwrap();
+//! # encode_request_header_into_buffer(&mut buf, &req_header).unwrap();
 //! # let mut api_versions_req = ApiVersionsRequest::default();
 //! # api_versions_req.client_software_version = StrBytes::from_static_str("1.0");
 //! # api_versions_req.client_software_name = StrBytes::from_static_str("example-client");
-//! # api_versions_req.encode(&mut buf, 3);
+//! # api_versions_req.encode(&mut buf, 3).unwrap();
 //!
-//! let api_key = buf.peek_bytes(0..2).get_i16();
-//! let api_version = buf.peek_bytes(2..4).get_i16();
-//! let header_version = ApiKey::try_from(api_key).unwrap().request_header_version(api_version);
 //!
-//! let header = RequestHeader::decode(&mut buf, header_version).unwrap();
+//! let header = decode_request_header_from_buffer(&mut buf).unwrap();
 //! let api_key = ApiKey::try_from(header.request_api_version);
 //! let req = match api_key {
 //!     ApiVersionsKey => RequestKind::ApiVersions(ApiVersionsRequest::decode(&mut buf, header.request_api_version).unwrap()),

src/protocol/mod.rs

@@ -121,6 +121,8 @@ mod str_bytes {
 
 pub use str_bytes::StrBytes;
 
+use crate::messages::{ApiKey, RequestHeader};
+
 pub(crate) trait NewType<Inner>: From<Inner> + Into<Inner> + Borrow<Inner> {}
 
 impl<T> NewType<T> for T {}
@@ -209,6 +211,26 @@ pub trait Request: Message + Encodable + Decodable + HeaderVersion {
     type Response: Message + Encodable + Decodable + HeaderVersion;
 }
 
+/// Decode the request header from the provided buffer.
+pub fn decode_request_header_from_buffer<B: ByteBuf>(buf: &mut B) -> Result<RequestHeader> {
+    let api_key = ApiKey::try_from(bytes::Buf::get_i16(&mut buf.peek_bytes(0..2)))
+        .map_err(|_| anyhow::Error::msg("Unknown API key"))?;
+    let api_version = bytes::Buf::get_i16(&mut buf.peek_bytes(2..4));
+    let header_version = api_key.request_header_version(api_version);
+    RequestHeader::decode(buf, header_version)
+}
+
+/// Encode the request header into the provided buffer.
+pub fn encode_request_header_into_buffer<B: ByteBufMut>(
+    buf: &mut B,
+    header: &RequestHeader,
+) -> Result<()> {
+    let api_key = ApiKey::try_from(header.request_api_key)
+        .map_err(|_| anyhow::Error::msg("Unknown API key"))?;
+    let version = api_key.request_header_version(header.request_api_version);
+    header.encode(buf, version)
+}
+
 pub(crate) fn write_unknown_tagged_fields<B: ByteBufMut, R: RangeBounds<i32>>(
     buf: &mut B,
     range: R,

tests/all_tests/common.rs

@@ -12,6 +12,7 @@ use testcontainers::{
     runners::SyncRunner,
     Container, GenericImage, ImageExt,
 };
+use kafka_protocol::protocol::encode_request_header_into_buffer;
 
 pub fn start_kafka() -> Container<GenericImage> {
     GenericImage::new("bitnami/kafka", "3.6.1-debian-11-r24")
@@ -48,9 +49,7 @@ pub fn send_request<T: Encodable + HeaderVersion>(
 ) {
     let mut bytes = BytesMut::new();
 
-    header
-        .encode(&mut bytes, T::header_version(header.request_api_version))
-        .unwrap();
+    encode_request_header_into_buffer(&mut bytes, &header).unwrap();
     body.encode(&mut bytes, header.request_api_version).unwrap();
 
     let size = bytes.len() as i32;

tests/all_tests/request_header.rs

@@ -1,22 +1,14 @@
-use bytes::{Buf, Bytes};
-use kafka_protocol::messages::{ApiKey, RequestHeader};
-use kafka_protocol::protocol::buf::ByteBuf;
-use kafka_protocol::protocol::Decodable;
-use std::convert::TryFrom;
+use bytes::Bytes;
+use kafka_protocol::protocol::decode_request_header_from_buffer;
 
 #[test]
 fn request_header() {
     let mut bytes = Bytes::from(vec![
         0x00, 0x12, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0d, 0x61, 0x64, 0x6d, 0x69, 0x6e,
         0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x2d, 0x31, 0x00,
     ]);
-
-    let api_key = bytes.peek_bytes(0..2).get_i16();
-    let api_version = bytes.peek_bytes(2..4).get_i16();
-    let header_version = ApiKey::try_from(api_key)
-        .unwrap()
-        .request_header_version(api_version);
-    let res = RequestHeader::decode(&mut bytes, header_version).unwrap();
+    
+    let res = decode_request_header_from_buffer(&mut bytes).unwrap();
 
     assert_eq!(res.request_api_key, 18);
     assert_eq!(res.request_api_version, 3);