From f55bd240f1496d28c553770e027802529222e120 Mon Sep 17 00:00:00 2001
From: umgefahren <55623006+umgefahren@users.noreply.github.com>
Date: Sat, 3 Feb 2024 19:08:37 +0100
Subject: [PATCH] Make no_std with alloc compatible

---
 Cargo.toml                      | 3 ++-
 src/builder.rs                  | 7 +++++--
 src/cipherstate.rs              | 3 +++
 src/error.rs                    | 3 ++-
 src/handshakestate.rs           | 6 +++++-
 src/lib.rs                      | 5 +++++
 src/params/mod.rs               | 6 +++++-
 src/params/patterns.rs          | 5 ++++-
 src/resolvers/default.rs        | 3 +++
 src/resolvers/mod.rs            | 3 +++
 src/resolvers/ring.rs           | 3 +++
 src/stateless_transportstate.rs | 2 +-
 src/symmetricstate.rs           | 3 +++
 src/transportstate.rs           | 2 +-
 src/utils.rs                    | 2 +-
 15 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 46060dcb..43d1278e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,7 +16,7 @@ edition = "2021"
 # Features with a -resolver suffix simply enables the existence of a specific resolver,
 # and -accelerated suffix means that this resolver will be the default used by the Builder.
 [features]
-default = ["default-resolver"]
+default = ["default-resolver", "std"]
 default-resolver = ["aes-gcm", "chacha20poly1305", "blake2", "sha2", "curve25519-dalek"]
 nightly = ["blake2/simd_opt", "subtle/nightly"]
 ring-resolver = ["ring"]
@@ -28,6 +28,7 @@ hfs = []
 pqclean_kyber1024 = ["pqcrypto-kyber", "pqcrypto-traits", "hfs", "default-resolver"]
 xchachapoly = ["chacha20poly1305", "default-resolver"]
 risky-raw-split = []
+std = []
 
 [[bench]]
 name = "benches"
diff --git a/src/builder.rs b/src/builder.rs
index 07f8aafa..546dcd0b 100644
--- a/src/builder.rs
+++ b/src/builder.rs
@@ -1,4 +1,4 @@
-use std::fmt::Debug;
+use core::fmt::Debug;
 
 #[cfg(feature = "hfs")]
 use crate::params::HandshakeModifier;
@@ -13,6 +13,9 @@ use crate::{
 };
 use subtle::ConstantTimeEq;
 
+#[cfg(not(feature = "std"))]
+use alloc::{boxed::Box, vec, vec::Vec};
+
 /// The maximum number of PSKs we will allocate for.
 const MAX_PSKS: usize = 10;
 
@@ -65,7 +68,7 @@ pub struct Builder<'builder> {
 }
 
 impl<'builder> Debug for Builder<'builder> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("Builder").field("params", &self.params.name).finish_non_exhaustive()
     }
 }
diff --git a/src/cipherstate.rs b/src/cipherstate.rs
index fa29dc45..4399d85d 100644
--- a/src/cipherstate.rs
+++ b/src/cipherstate.rs
@@ -4,6 +4,9 @@ use crate::{
     types::Cipher,
 };
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
 pub(crate) struct CipherState {
     cipher:  Box<dyn Cipher>,
     n:       u64,
diff --git a/src/error.rs b/src/error.rs
index 8734ec4f..acda40f5 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,6 +1,6 @@
 //! All error types used by Snow operations.
 
-use std::fmt;
+use core::fmt;
 
 /// `snow` provides decently detailed errors, exposed as the [`Error`] enum,
 /// to allow developers to react to errors in a more actionable way.
@@ -181,4 +181,5 @@ impl fmt::Display for Error {
     }
 }
 
+#[cfg(feature = "std")]
 impl std::error::Error for Error {}
diff --git a/src/handshakestate.rs b/src/handshakestate.rs
index 671b55dc..ccec459a 100644
--- a/src/handshakestate.rs
+++ b/src/handshakestate.rs
@@ -15,11 +15,15 @@ use crate::{
     types::{Dh, Hash, Random},
     utils::Toggle,
 };
-use std::{
+use core::{
     convert::{TryFrom, TryInto},
     fmt,
 };
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
+
 /// A state machine encompassing the handshake phase of a Noise session.
 ///
 /// **Note:** you are probably looking for [`Builder`](struct.Builder.html) to
diff --git a/src/lib.rs b/src/lib.rs
index 53365dc6..7a1777c8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -53,8 +53,13 @@
 //!
 //! See `examples/simple.rs` for a more complete TCP client/server example with static keys.
 
+#![cfg_attr(not(feature = "std"), no_std)]
 #![warn(missing_docs)]
 
+
+#[cfg(not(feature = "std"))]
+extern crate alloc;
+
 macro_rules! copy_slices {
     ($inslice:expr, $outslice:expr) => {
         $outslice[..$inslice.len()].copy_from_slice(&$inslice[..])
diff --git a/src/params/mod.rs b/src/params/mod.rs
index f8d7463d..693af2eb 100644
--- a/src/params/mod.rs
+++ b/src/params/mod.rs
@@ -5,9 +5,13 @@
 //! patterns/names)
 
 use crate::error::{Error, PatternProblem};
-use std::str::FromStr;
+use core::str::FromStr;
 mod patterns;
 
+#[cfg(not(feature = "std"))]
+use alloc::{string::String, borrow::ToOwned};
+
+
 pub use self::patterns::{
     HandshakeChoice, HandshakeModifier, HandshakeModifierList, HandshakePattern,
     SUPPORTED_HANDSHAKE_PATTERNS,
diff --git a/src/params/patterns.rs b/src/params/patterns.rs
index ffff2414..9babffc4 100644
--- a/src/params/patterns.rs
+++ b/src/params/patterns.rs
@@ -1,7 +1,10 @@
 #![allow(clippy::enum_glob_use)]
 
 use crate::error::{Error, PatternProblem};
-use std::{convert::TryFrom, str::FromStr};
+use core::{convert::TryFrom, str::FromStr};
+
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
 
 /// A small helper macro that behaves similar to the `vec![]` standard macro,
 /// except it allocates a bit extra to avoid resizing.
diff --git a/src/resolvers/default.rs b/src/resolvers/default.rs
index 93ed5b6a..d3dde829 100644
--- a/src/resolvers/default.rs
+++ b/src/resolvers/default.rs
@@ -22,6 +22,9 @@ use crate::{
     Error,
 };
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
 /// The default resolver provided by snow. This resolver is designed to
 /// support as many of the Noise spec primitives as possible with
 /// pure-Rust (or nearly pure-Rust) implementations.
diff --git a/src/resolvers/mod.rs b/src/resolvers/mod.rs
index e6c15dd5..a9274a71 100644
--- a/src/resolvers/mod.rs
+++ b/src/resolvers/mod.rs
@@ -26,6 +26,9 @@ pub use self::libsodium::SodiumResolver;
 #[cfg(feature = "ring-resolver")]
 pub use self::ring::RingResolver;
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
 /// Boxed `CryptoResolver`
 pub type BoxedCryptoResolver = Box<dyn CryptoResolver + Send>;
 
diff --git a/src/resolvers/ring.rs b/src/resolvers/ring.rs
index 83180fbb..791c3a4f 100644
--- a/src/resolvers/ring.rs
+++ b/src/resolvers/ring.rs
@@ -11,6 +11,9 @@ use ring::{
     rand::{SecureRandom, SystemRandom},
 };
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
 /// A resolver that chooses [ring](https://github.com/briansmith/ring)-backed
 /// primitives when available.
 #[allow(clippy::module_name_repetitions)]
diff --git a/src/stateless_transportstate.rs b/src/stateless_transportstate.rs
index 904ab4a5..1fd52014 100644
--- a/src/stateless_transportstate.rs
+++ b/src/stateless_transportstate.rs
@@ -6,7 +6,7 @@ use crate::{
     params::HandshakePattern,
     utils::Toggle,
 };
-use std::{convert::TryFrom, fmt};
+use core::{convert::TryFrom, fmt};
 
 /// A state machine encompassing the transport phase of a Noise session, using the two
 /// `CipherState`s (for sending and receiving) that were spawned from the `SymmetricState`'s
diff --git a/src/symmetricstate.rs b/src/symmetricstate.rs
index 7114cd30..6c1fc638 100644
--- a/src/symmetricstate.rs
+++ b/src/symmetricstate.rs
@@ -5,6 +5,9 @@ use crate::{
     types::Hash,
 };
 
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
+
 #[derive(Copy, Clone)]
 pub(crate) struct SymmetricStateData {
     h:       [u8; MAXHASHLEN],
diff --git a/src/transportstate.rs b/src/transportstate.rs
index b927b37f..47e79435 100644
--- a/src/transportstate.rs
+++ b/src/transportstate.rs
@@ -6,7 +6,7 @@ use crate::{
     params::HandshakePattern,
     utils::Toggle,
 };
-use std::{convert::TryFrom, fmt};
+use core::{convert::TryFrom, fmt};
 
 /// A state machine encompassing the transport phase of a Noise session, using the two
 /// `CipherState`s (for sending and receiving) that were spawned from the `SymmetricState`'s
diff --git a/src/utils.rs b/src/utils.rs
index 633bea94..4e6efb47 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,4 +1,4 @@
-use std::ops::{Deref, DerefMut};
+use core::ops::{Deref, DerefMut};
 
 /// Toggle is similar to Option, except that even in the Off/"None" case, there is still
 /// an owned allocated inner object. This is useful for holding onto pre-allocated objects