Support big endian headers in addition to little endian.

fizyr · Dec 4, 2023 · 58b14dc · 58b14dc
1 parent b568d9b
commit 58b14dc
Show file tree

Hide file tree

Showing 10 changed files with 123 additions and 57 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,8 @@
+# Unreleased
+- [change][major] Mark `StreamConfig` and `UnixConfig` as non-exhaustive structs.
+- [change][major] Make the `MessageHeader::encode/decode()` functions take an `endian` parameter.
+- [add][major] Add an `endian` field to `StreamConfig` and `UnixConfig`.
+
 # Version 0.7.1 - 2023-11-26
 - [change][patch] Remove dependency on `byteorder` crate.
 

diff --git a/src/message.rs b/src/message.rs
@@ -1,5 +1,6 @@
 use crate::Error;
 use crate::error::private::InnerError;
+use crate::transport::Endian;
 
 /// The encoded length of a message header.
 ///
@@ -220,16 +221,16 @@ impl MessageHeader {
 		}
 	}
 
-	/// Decode a message header from a byte slice.
+	/// Decode a message header from a byte slice using the given endianness for the header fields.
 	///
 	/// The byte slice should NOT contain the message size.
 	///
 	/// # Panic
 	/// This function panics if the buffer does not contain a full header.
-	pub fn decode(buffer: &[u8]) -> Result<Self, Error> {
-		let message_type = read_u32_le(&buffer[0..]);
-		let request_id = read_u32_le(&buffer[4..]);
-		let service_id = read_i32_le(&buffer[8..]);
+	pub fn decode(buffer: &[u8], endian: Endian) -> Result<Self, Error> {
+		let message_type = endian.read_u32(&buffer[0..]);
+		let request_id = endian.read_u32(&buffer[4..]);
+		let service_id = endian.read_i32(&buffer[8..]);
 
 		let message_type = MessageType::from_u32(message_type)?;
 		Ok(Self {
@@ -239,17 +240,17 @@ impl MessageHeader {
 		})
 	}
 
-	/// Encode a message header into a byte slice.
+	/// Encode a message header into a byte slice using the given endianness for the header fields.
 	///
 	/// This will NOT add a message size (which would be impossible even if we wanted to).
 	///
 	/// # Panic
 	/// This function panics if the buffer is not large enough to hold a full header.
-	pub fn encode(&self, buffer: &mut [u8]) {
+	pub fn encode(&self, buffer: &mut [u8], endian: Endian) {
 		assert!(buffer.len() >= 12);
-		write_u32_le(&mut buffer[0..], self.message_type as u32);
-		write_u32_le(&mut buffer[4..], self.request_id);
-		write_i32_le(&mut buffer[8..], self.service_id);
+		endian.write_u32(&mut buffer[0..], self.message_type as u32);
+		endian.write_u32(&mut buffer[4..], self.request_id);
+		endian.write_i32(&mut buffer[8..], self.service_id);
 	}
 }
 
@@ -260,23 +261,3 @@ impl<Body> std::fmt::Debug for Message<Body> {
 			.finish_non_exhaustive()
 	}
 }
-
-/// Read a [`u32`] from a buffer in little endian format.
-fn read_u32_le(buffer: &[u8]) -> u32 {
-	u32::from_le_bytes(buffer[0..4].try_into().unwrap())
-}
-
-/// Read a [`i32`] from a buffer in little endian format.
-fn read_i32_le(buffer: &[u8]) -> i32 {
-	i32::from_le_bytes(buffer[0..4].try_into().unwrap())
-}
-
-/// Write a [`i32`] to a buffer in little endian format.
-fn write_i32_le(buffer: &mut [u8], value: i32) {
-	buffer[0..4].copy_from_slice(&value.to_le_bytes());
-}
-
-/// Write a [`u32`] to a buffer in little endian format.
-fn write_u32_le(buffer: &mut [u8], value: u32) {
-	buffer[0..4].copy_from_slice(&value.to_le_bytes());
-}
diff --git a/src/transport/endian.rs b/src/transport/endian.rs
@@ -0,0 +1,50 @@
+/// The endianness to use for encoding header fields.
+///
+/// The encoding and serialization of message bodies is up to the application code,
+/// and it not affected by this configuration parameter.
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
+pub enum Endian {
+	/// Encode header fields in little endian.
+	LittleEndian,
+
+	/// Encode header fields in big endian.
+	BigEndian,
+}
+
+impl Endian {
+	/// Read a [`u32`] from a buffer in the correct endianness.
+	pub(crate) fn read_u32(self, buffer: &[u8]) -> u32 {
+		let buffer = buffer[0..4].try_into().unwrap();
+		match self {
+			Self::LittleEndian => u32::from_le_bytes(buffer),
+			Self::BigEndian => u32::from_be_bytes(buffer),
+		}
+	}
+
+	/// Write a [`u32`] to a buffer in thcorrect endianness.
+	pub(crate) fn write_u32(self, buffer: &mut [u8], value: u32) {
+		let bytes = match self {
+			Self::LittleEndian => value.to_le_bytes(),
+			Self::BigEndian => value.to_be_bytes(),
+		};
+		buffer[0..4].copy_from_slice(&bytes);
+	}
+
+	/// Read a [`i32`] from a buffer in the correct endianness.
+	pub(crate) fn read_i32(self, buffer: &[u8]) -> i32 {
+		let buffer = buffer[0..4].try_into().unwrap();
+		match self {
+			Self::LittleEndian => i32::from_le_bytes(buffer),
+			Self::BigEndian => i32::from_be_bytes(buffer),
+		}
+	}
+
+	/// Write a [`i32`] to a buffer in thcorrect endianness.
+	pub(crate) fn write_i32(self, buffer: &mut [u8], value: i32) {
+		let bytes = match self {
+			Self::LittleEndian => value.to_le_bytes(),
+			Self::BigEndian => value.to_be_bytes(),
+		};
+		buffer[0..4].copy_from_slice(&bytes);
+	}
+}
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
@@ -12,6 +12,9 @@ use std::task::{Context, Poll};
 
 use crate::{Error, Message, MessageHeader};
 
+mod endian;
+pub use endian::Endian;
+
 pub(crate) mod stream;
 pub use stream::StreamTransport;
 

diff --git a/src/transport/stream/config.rs b/src/transport/stream/config.rs
@@ -1,5 +1,8 @@
+use crate::transport::Endian;
+
 /// Configuration for a byte-stream transport.
 #[derive(Debug, Clone)]
+#[non_exhaustive]
 pub struct StreamConfig {
 	/// The maximum body size for incoming messages.
 	///
@@ -13,13 +16,20 @@ pub struct StreamConfig {
 	/// the message is discarded and an error is returned.
 	/// Stream sockets remain usable since the message header will not be sent either.
 	pub max_body_len_write: u32,
+
+	/// The endianness to use when encoding/decoding header fields.
+	///
+	/// The encoding and serialization of message bodies is up to the application code,
+	/// and it not affected by this configuration parameter.
+	pub endian: Endian,
 }
 
 impl Default for StreamConfig {
 	fn default() -> Self {
 		Self {
 			max_body_len_read: 8 * 1024,
 			max_body_len_write: 8 * 1024,
+			endian: Endian::LittleEndian,
 		}
 	}
 }
diff --git a/src/transport/stream/mod.rs b/src/transport/stream/mod.rs
@@ -53,8 +53,8 @@ mod impl_unix_stream {
 
 		fn split(&mut self) -> (StreamReadHalf<tokio::net::unix::ReadHalf>, StreamWriteHalf<tokio::net::unix::WriteHalf>) {
 			let (read_half, write_half) = self.stream.split();
-			let read_half = StreamReadHalf::new(read_half, self.config.max_body_len_read);
-			let write_half = StreamWriteHalf::new(write_half, self.config.max_body_len_write);
+			let read_half = StreamReadHalf::new(read_half, self.config.max_body_len_read, self.config.endian);
+			let write_half = StreamWriteHalf::new(write_half, self.config.max_body_len_write, self.config.endian);
 			(read_half, write_half)
 		}
 
@@ -155,8 +155,8 @@ mod impl_tcp {
 
 		fn split(&mut self) -> (StreamReadHalf<tokio::net::tcp::ReadHalf>, StreamWriteHalf<tokio::net::tcp::WriteHalf>) {
 			let (read_half, write_half) = self.stream.split();
-			let read_half = StreamReadHalf::new(read_half, self.config.max_body_len_read);
-			let write_half = StreamWriteHalf::new(write_half, self.config.max_body_len_write);
+			let read_half = StreamReadHalf::new(read_half, self.config.max_body_len_read, self.config.endian);
+			let write_half = StreamWriteHalf::new(write_half, self.config.max_body_len_write, self.config.endian);
 			(read_half, write_half)
 		}
 

diff --git a/src/transport/stream/transport.rs b/src/transport/stream/transport.rs
@@ -5,7 +5,7 @@ use tokio::io::{AsyncRead, AsyncWrite};
 
 use super::{StreamBody, StreamConfig};
 use crate::error::private::check_payload_too_large;
-use crate::transport::TransportError;
+use crate::transport::{TransportError, Endian};
 use crate::{Message, MessageHeader};
 
 /// Length of a message frame and header.
@@ -30,6 +30,9 @@ pub struct StreamReadHalf<ReadStream> {
 	/// The maximum body length to accept when reading messages.
 	pub(super) max_body_len: u32,
 
+	/// The endianness to use for decoding header fields.
+	pub(super) endian: Endian,
+
 	/// The number of bytes read for the current message.
 	pub(super) bytes_read: usize,
 
@@ -52,6 +55,9 @@ pub struct StreamWriteHalf<WriteStream> {
 	/// The maximum body length to enforce for messages.
 	pub(super) max_body_len: u32,
 
+	/// The endianness to use for encoding header fields.
+	pub(super) endian: Endian,
+
 	/// The number of bytes written for the current message.
 	pub(super) bytes_written: usize,
 
@@ -91,10 +97,11 @@ where
 
 impl<ReadStream> StreamReadHalf<ReadStream> {
 	#[allow(dead_code)] // Not used when transports are disabled.
-	pub(super) fn new(stream: ReadStream, max_body_len: u32) -> Self {
+	pub(super) fn new(stream: ReadStream, max_body_len: u32, endian: Endian) -> Self {
 		Self {
 			stream,
 			max_body_len,
+			endian,
 			header_buffer: [0u8; FRAMED_HEADER_LEN],
 			bytes_read: 0,
 			parsed_header: MessageHeader::request(0, 0),
@@ -117,10 +124,11 @@ impl<ReadStream> StreamReadHalf<ReadStream> {
 
 impl<WriteStream> StreamWriteHalf<WriteStream> {
 	#[allow(dead_code)] // Not used when transports are disabled.
-	pub(super) fn new(stream: WriteStream, max_body_len: u32) -> Self {
+	pub(super) fn new(stream: WriteStream, max_body_len: u32, endian: Endian) -> Self {
 		Self {
 			stream,
 			max_body_len,
+			endian,
 			header_buffer: None,
 			bytes_written: 0,
 		}
@@ -171,8 +179,8 @@ where
 			// Check if we have the whole frame + header.
 			if this.bytes_read == FRAMED_HEADER_LEN {
 				// Parse frame and header.
-				let length = read_u32_le(&this.header_buffer[0..]);
-				this.parsed_header = MessageHeader::decode(&this.header_buffer[4..])
+				let length = this.endian.read_u32(&this.header_buffer[0..]);
+				this.parsed_header = MessageHeader::decode(&this.header_buffer[4..], this.endian)
 					.map_err(TransportError::new_fatal)?;
 
 				// Check body length and create body buffer.
@@ -218,8 +226,8 @@ where
 		// Encode the header if we haven't done that yet.
 		let header_buffer = this.header_buffer.get_or_insert_with(|| {
 			let mut buffer = [0u8; FRAMED_HEADER_LEN];
-			write_u32_le(&mut buffer[0..], body.len() as u32 + crate::HEADER_LEN);
-			header.encode(&mut buffer[4..]);
+			this.endian.write_u32(&mut buffer[0..], body.len() as u32 + crate::HEADER_LEN);
+			header.encode(&mut buffer[4..], this.endian);
 			buffer
 		});
 
@@ -241,13 +249,3 @@ where
 		Poll::Ready(Ok(()))
 	}
 }
-
-/// Read a [`u32`] from a buffer in little endian format.
-fn read_u32_le(buffer: &[u8]) -> u32 {
-	u32::from_le_bytes(buffer[0..4].try_into().unwrap())
-}
-
-/// Write a [`u32`] to a buffer in little endian format.
-fn write_u32_le(buffer: &mut [u8], value: u32) {
-	buffer[0..4].copy_from_slice(&value.to_le_bytes());
-}
diff --git a/src/transport/unix/config.rs b/src/transport/unix/config.rs
@@ -1,5 +1,8 @@
+use crate::transport::Endian;
+
 /// Configuration for Unix datagram transports.
 #[derive(Debug, Clone)]
+#[non_exhaustive]
 pub struct UnixConfig {
 	/// The maximum body size for incoming messages.
 	///
@@ -22,6 +25,12 @@ pub struct UnixConfig {
 
 	/// The maximum number of attached file descriptors for sending messages.
 	pub max_fds_write: u32,
+
+	/// The endianness to use when encoding/decoding header fields.
+	///
+	/// The encoding and serialization of message bodies is up to the application code,
+	/// and it not affected by this configuration parameter.
+	pub endian: Endian,
 }
 
 impl Default for UnixConfig {
@@ -31,6 +40,7 @@ impl Default for UnixConfig {
 			max_body_len_write: 4 * 1024,
 			max_fds_read: 10,
 			max_fds_write: 10,
+			endian: Endian::LittleEndian,
 		}
 	}
 }
diff --git a/src/transport/unix/mod.rs b/src/transport/unix/mod.rs
@@ -53,8 +53,8 @@ mod impl_unix_seqpacket {
 
 		fn split(&mut self) -> (UnixReadHalf<&tokio_seqpacket::UnixSeqpacket>, UnixWriteHalf<&tokio_seqpacket::UnixSeqpacket>) {
 			let (read_half, write_half) = (&self.socket, &self.socket);
-			let read_half = UnixReadHalf::new(read_half, self.config.max_body_len_read, self.config.max_fds_read);
-			let write_half = UnixWriteHalf::new(write_half, self.config.max_body_len_write, self.config.max_fds_write);
+			let read_half = UnixReadHalf::new(read_half, self.config.max_body_len_read, self.config.max_fds_read, self.config.endian);
+			let write_half = UnixWriteHalf::new(write_half, self.config.max_body_len_write, self.config.max_fds_write, self.config.endian);
 			(read_half, write_half)
 		}