From f2514afac19d705c67de6b21e11673723c13b4d2 Mon Sep 17 00:00:00 2001
From: Kiran Shila <me@kiranshila.com>
Date: Fri, 5 Apr 2024 15:44:56 -0700
Subject: [PATCH] Fix endianness check bug

---
 Cargo.toml   |  2 +-
 src/read.rs  | 42 +++++++++++++++++++++++++++++-------------
 src/write.rs | 29 +++++++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 34041e2..91ef40c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "sigproc_filterbank"
-version = "0.3.1"
+version = "0.4.0"
 edition = "2021"
 rust-version = "1.57.0"
 license = "Apache-2.0 OR MIT"
diff --git a/src/read.rs b/src/read.rs
index 7184f1b..b88bfcb 100644
--- a/src/read.rs
+++ b/src/read.rs
@@ -18,6 +18,16 @@ use crate::errors::FilterbankError;
 type ParseResult<'a, T> = IResult<&'a [u8], T, FilterbankError>;
 type HeaderResult<'a> = ParseResult<'a, HeaderParameter<'a>>;
 
+/// Figure out the endianness of the file by matching the length field of the `HEADER_START` at the start of the file
+fn determine_endianness(input: &[u8]) -> Result<Endianness, nom::Err<FilterbankError>> {
+    let test_bytes = &input[..4];
+    match test_bytes {
+        [12, 0, 0, 0] => Ok(Endianness::Little),
+        [0, 0, 0, 12] => Ok(Endianness::Big),
+        _ => Err(nom::Err::Failure(FilterbankError::InvalidHeader)),
+    }
+}
+
 fn header_string(s: &'static str, endian: Endianness) -> impl FnMut(&[u8]) -> ParseResult<&[u8]> {
     move |input: &[u8]| length_value(u32(endian), tag(s))(input)
 }
@@ -174,16 +184,10 @@ fn nbits(input: &[u8], endian: Endianness) -> HeaderResult {
 }
 
 fn header<'a>(input: &'a [u8]) -> ParseResult<'a, (Endianness, Vec<HeaderParameter<'a>>)> {
-    // Determine the endianness based on the first match of HEADER_START
-    let res_big = header_start(input, Endianness::Big);
-    let res_little = header_start(input, Endianness::Little);
-    let (remaining, endian) = if let Ok((remaining, _)) = res_big {
-        (remaining, Endianness::Big)
-    } else if let Ok((remaining, _)) = res_little {
-        (remaining, Endianness::Little)
-    } else {
-        return Err(res_little.err().unwrap());
-    };
+    // Determine the endianness
+    let endian = determine_endianness(input)?;
+    // Match the header
+    let (remaining, _) = header_start(input, endian)?;
     // The rest of the owl
     let (remaining, (headers, _)) = many_till(
         alt((
@@ -511,7 +515,7 @@ mod tests {
     use std::{fs::File, io::Read};
 
     use super::*;
-    use crate::write::sigproc_string;
+    use crate::write::{sigproc_string, sigproc_string_endian};
 
     #[test]
     fn test_start_end() {
@@ -524,8 +528,20 @@ mod tests {
     }
 
     #[test]
-    fn test_wrong_endian() {
-        let hstart = sigproc_string("HEADER_START");
+    fn test_wrong_endian_big() {
+        let hstart = sigproc_string_endian("HEADER_START", Endianness::Big);
+        let a = header_start(&hstart, Endianness::Big);
+        let b = header_start(&hstart, Endianness::Little);
+        if a.is_err() {
+            assert!(b.is_ok())
+        } else {
+            assert!(b.is_err())
+        }
+    }
+
+    #[test]
+    fn test_wrong_endian_little() {
+        let hstart = sigproc_string_endian("HEADER_START", Endianness::Little);
         let a = header_start(&hstart, Endianness::Big);
         let b = header_start(&hstart, Endianness::Little);
         if a.is_err() {
diff --git a/src/write.rs b/src/write.rs
index 0d83ee9..c79ee0f 100644
--- a/src/write.rs
+++ b/src/write.rs
@@ -1,7 +1,22 @@
 use std::marker::PhantomData;
 
+use nom::number::Endianness;
 use ux::{u1, u2, u4};
 
+/// Create an endian-specific sigproc-compatible string
+#[allow(unused)]
+pub(crate) fn sigproc_string_endian(s: &str, endian: Endianness) -> Vec<u8> {
+    let len = s.len() as u32;
+    let mut out = vec![];
+    match endian {
+        Endianness::Big => out.extend_from_slice(&len.to_be_bytes()),
+        Endianness::Little => out.extend_from_slice(&len.to_le_bytes()),
+        Endianness::Native => out.extend_from_slice(&len.to_ne_bytes()),
+    }
+    out.extend_from_slice(s.as_bytes());
+    out
+}
+
 /// Creates a sigproc-compatible string
 pub(crate) fn sigproc_string(s: &str) -> Vec<u8> {
     let len = s.len() as u32;
@@ -313,11 +328,21 @@ mod tests {
     use super::*;
     use crate::read::ReadFilterbank;
 
+    // Tested seperately because we don't know the endianness of the test platform
+
     #[test]
-    fn test_sigproc_string() {
+    fn test_sigproc_string_little() {
         assert_eq!(
             b"\x0C\x00\x00\x00HEADER_START".to_vec(),
-            sigproc_string("HEADER_START")
+            sigproc_string_endian("HEADER_START", Endianness::Little)
+        );
+    }
+
+    #[test]
+    fn test_sigproc_string_big() {
+        assert_eq!(
+            b"\x00\x00\x00\x0CHEADER_START".to_vec(),
+            sigproc_string_endian("HEADER_START", Endianness::Big)
         );
     }