Add no_std support (#13)

* Replace imports from std:: with core:: * Move 'clap' as a dependency for the 'std' feature Introduce a new `std` feature that can be optionally disabled for nostd support. The new feature is on by default, so the default behaviour of this crate remains unchanged. Note that building with `nostd` does not yet work. * Add a prelude to import common features from core Some components like `TryInto`, `String, etc are not visible by default with nostd (since they are exposed via std's prelude). Create a crate-local prelude to expose all these. * Use a custom error type Drops the dependency on std::io::Error. This finally enables building with `nostd`, which can be tested via: cargo build --target x86_64-unknown-uefi --no-default-features --lib Note that trying to build the binary (e.g.: `main.rs`) with nostd is not supported. Co-Authored-By: Isaac Marovitz <[email protected]> --------- Co-authored-by: Isaac Marovitz <[email protected]>
IsaacMarovitz · May 27, 2023 · 9c88a9e · 9c88a9e
1 parent 3412036
commit 9c88a9e
Show file tree

Hide file tree

Showing 8 changed files with 87 additions and 30 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -12,13 +12,17 @@ categories = ["command-line-utilities"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+default = ["std"]
+std = ["clap"]
+
 [dependencies]
 bytemuck = {version = "1.13.1", features = ["derive"]}
 num-traits = { version = "0.2", default-features = false }
 num-derive = "0.3"
 bitflags = { version = "2.3.1", default-features = false }
 chrono = { version = "0.4.24", default-features = false }
-clap = { version = "4.3.0", features = ["cargo"] }
+clap = { version = "4.3.0", features = ["cargo"], optional = true }
 
 [dev-dependencies]
 datatest-stable = "0.1.3"

diff --git a/src/coff.rs b/src/coff.rs
@@ -2,8 +2,9 @@ use bytemuck::{Pod, Zeroable};
 use num_derive::FromPrimitive;   
 use num_traits::FromPrimitive;
 use bitflags::bitflags;
-use std::{fmt, str};
+use core::{fmt, str};
 use chrono::NaiveDateTime;
+use crate::prelude::*;
 
 /// COFF File Header (Object and Image)
 #[derive(Copy, Clone, Pod, Zeroable, Default)]

diff --git a/src/lib.rs b/src/lib.rs
@@ -10,7 +10,7 @@
 //! # use std::{fs, io};
 //! use pe_parser::pe::parse_portable_executable;
 //! 
-//! # fn main() -> io::Result<()> {
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
 //! # let path_to_pe = "tests/pe/64_pe/64_pe_checksum_non_zero.dat";
 //! // Read the binary from a file
 //! let binary = fs::read(path_to_pe)?;
@@ -24,6 +24,11 @@
 //! ```
 
 #![warn(missing_docs)]
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+use crate::prelude::*;
+use core::fmt;
 
 /// COFF file header definitions and helper functions
 pub mod coff;
@@ -33,4 +38,38 @@ pub mod optional;
 pub mod section;
 /// Monolith struct containing all the information
 /// you will ever need
-pub mod pe;
+pub mod pe;
+mod prelude;
+
+/// Error parsing a PE binary.
+#[derive(Debug)]
+pub enum Error {
+    /// Failed to read data; premature EOF.
+    OffsetOutOfRange,
+    /// Failed to parse a header for an optional.
+    BadOptionalHeader,
+    /// Failed to parse a String.
+    BadString(alloc::string::FromUtf8Error),
+    /// Missing PE header.
+    MissingPeHeader,
+    /// Missing COFF header.
+    MissingCoffHeader,
+    /// Missing magic number from header.
+    MissingMagicNumber,
+}
+
+impl core::fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Error::OffsetOutOfRange => f.write_str("Offset out of range!"),
+            Error::BadOptionalHeader => f.write_str("Failed to parse optional header!"),
+            Error::BadString(e) => f.write_fmt(format_args!("Failed to parse string: {}!", e)),
+            Error::MissingPeHeader => f.write_str("Missing PE header!"),
+            Error::MissingCoffHeader => f.write_str("Missing COFF header!"),
+            Error::MissingMagicNumber => f.write_str("Missing magic number!"),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for Error {}
diff --git a/src/main.rs b/src/main.rs
@@ -1,9 +1,9 @@
-use std::io::{Error};
-use std::{env, fs};
+use core::env;
+use std::fs;
 use pe_parser::pe::parse_portable_executable;
 use clap::{Arg, command, ArgAction};
 
-fn main() -> Result<(), Error> {
+fn main() -> Result<(), Box<dyn std::error::Error>> {
     let matches = command!()
         .arg(Arg::new("file")
             .action(ArgAction::Set)

diff --git a/src/optional.rs b/src/optional.rs
@@ -2,8 +2,8 @@ use bytemuck::{Pod, Zeroable, checked::{try_from_bytes}};
 use num_derive::FromPrimitive;    
 use num_traits::FromPrimitive;
 use bitflags::bitflags;
-use std::{fmt, str};
-use std::io::{Error, ErrorKind};
+use core::{fmt, str};
+use crate::{prelude::*, Error};
 
 /// Magic values that determine if an Optional Header is 
 /// PE32 (32-bit) or PE32+ (64-bit)
@@ -444,11 +444,11 @@ impl Optional for optional_header_32 {
     }
 
     fn parse_optional_header(binary: &[u8], offset: &mut usize) -> Result<Self, Error> {
-        let size = std::mem::size_of::<Self>();
+        let size = core::mem::size_of::<Self>();
         let slice = match binary.get(*offset..*offset+size) {
             Some(slice) => slice,
             None => {
-                return Err(Error::new(ErrorKind::Other, "Offset out of range!"));
+                return Err(Error::OffsetOutOfRange);
             }
         };
 
@@ -460,7 +460,7 @@ impl Optional for optional_header_32 {
                 return Ok(header);
             }
             Err(_) => {
-                return Err(Error::new(ErrorKind::Other, "Failed to parse header!"));
+                return Err(Error::BadOptionalHeader);
             }
         }
     }
@@ -476,11 +476,11 @@ impl Optional for optional_header_64 {
     }
 
     fn parse_optional_header(binary: &[u8], offset: &mut usize) -> Result<Self, Error> {
-        let size = std::mem::size_of::<Self>();
+        let size = core::mem::size_of::<Self>();
         let slice = match binary.get(*offset..*offset+size) {
             Some(slice) => slice,
             None => {
-                return Err(Error::new(ErrorKind::Other, "Offset out of range!"));
+                return Err(Error::OffsetOutOfRange);
             }
         };
 
@@ -491,7 +491,7 @@ impl Optional for optional_header_64 {
                 return Ok(header);
             }
             Err(_) => {
-                return Err(Error::new(ErrorKind::Other, "Failed to parse header!"));
+                return Err(Error::BadOptionalHeader);
             }
         }
     }

diff --git a/src/pe.rs b/src/pe.rs
@@ -1,8 +1,8 @@
-use crate::{coff::coff_file_header, optional::{optional_header_32, optional_header_64, Magic, Optional}, section::{section_header, parse_section_table}};
-use std::io::{Error, ErrorKind};
+use crate::{coff::coff_file_header, optional::{optional_header_32, optional_header_64, Magic, Optional}, section::{section_header, parse_section_table}, Error};
 use bytemuck::checked::try_from_bytes;
 use num_traits::FromPrimitive;
-use std::fmt;
+use core::fmt;
+use crate::prelude::*;
 
 const IMAGE_DOS_PE_SIGNATURE_OFFSET: usize = 0x3c;
 
@@ -25,19 +25,19 @@ pub fn parse_portable_executable(binary: &[u8]) -> Result<PortableExecutable, Er
     let slice = match binary.get(offset..offset+4) {
         Some(slice) => slice,
         None => {
-            return Err(Error::new(ErrorKind::Other, "Offset out of range!"))
+            return Err(Error::OffsetOutOfRange);
         }
     };
 
     let string = match String::from_utf8(slice.to_vec()) {
         Ok(string) => string,
-        Err(_) => {
-            return Err(Error::new(ErrorKind::Other, "Failed to parse string!"));
+        Err(e) => {
+            return Err(Error::BadString(e));
         }
     };
 
     if string != "PE\0\0" {
-        return Err(Error::new(ErrorKind::InvalidData, "File is not a valid PE!"));
+        return Err(Error::MissingPeHeader);
     }
 
     offset += 4;
@@ -52,14 +52,14 @@ pub fn parse_portable_executable(binary: &[u8]) -> Result<PortableExecutable, Er
     let slice = match binary.get(offset..offset+20) {
         Some(slice) => slice,
         None => {
-            return Err(Error::new(ErrorKind::Other, "Offset out of range!"))
+            return Err(Error::OffsetOutOfRange);
         }
     };
 
     pe.coff = match try_from_bytes::<coff_file_header>(slice) {
         Ok(coff) => *coff,
         Err(_) => {
-            return Err(Error::new(ErrorKind::Other, "Failed to get COFF header!"));
+            return Err(Error::MissingCoffHeader);
         }
     };
 
@@ -69,7 +69,7 @@ pub fn parse_portable_executable(binary: &[u8]) -> Result<PortableExecutable, Er
         let magic = match Magic::from_u16(read_u16(binary, offset)?) {
             Some(magic) => magic,
             None => {
-                return Err(Error::new(ErrorKind::Other, "Failed to get Magic!"));
+                return Err(Error::MissingMagicNumber);
             }
         };
 
@@ -135,10 +135,13 @@ impl fmt::Display for PortableExecutable {
 
 fn read_u16(binary: &[u8], offset: usize) -> Result<u16, Error> {
     if let Some(array) = binary.get(offset..offset+2) {
-        if let Some(slice) = array.try_into().ok() {
-            return Ok(u16::from_le_bytes(slice));
+        if let Ok(slice) = array.try_into() {
+            Ok(u16::from_le_bytes(slice))
+        } else {
+            unreachable!()
         }
+    } else {
+        Err(Error::OffsetOutOfRange)
     }
 
-    Err(Error::new(ErrorKind::Other, "Failed to get value"))
 }
diff --git a/src/prelude.rs b/src/prelude.rs
@@ -0,0 +1,8 @@
+pub use core::prelude::v1::*;
+
+extern crate alloc;
+pub use alloc::string::String;
+pub use alloc::vec::Vec;
+pub use core::{write, writeln};
+pub use core::convert::TryInto;
+pub use core::unreachable;
diff --git a/src/section.rs b/src/section.rs
@@ -1,15 +1,17 @@
 use bytemuck::checked::try_from_bytes;
 use bytemuck::{Pod, Zeroable};
 use bitflags::bitflags;
-use std::{fmt, str};
+use core::{fmt, str};
+use core::writeln;
+use crate::prelude::*;
 
 /// Parse the section table from a byte array at a given offset.
 /// `number_of_sections` should be equal to number of sections
 /// defined in the COFF header.
 pub fn parse_section_table(binary: &[u8], offset: usize, number_of_sections: u16) -> Vec<section_header> {
     let mut offset = offset;
     let mut headers: Vec<section_header> = Vec::new();
-    let header_size = std::mem::size_of::<section_header>();
+    let header_size = core::mem::size_of::<section_header>();
 
     for _ in 0..number_of_sections {
         if let Some(slice) = binary.get(offset..offset+header_size) {