From e100b72c60d5a5b216e0743fcdb60ee89ad6436e Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Tue, 13 Aug 2024 10:48:00 -0600
Subject: [PATCH] ssh-encoding: add `Encode::{encode_vec, encode_bytes}`

Adds convenience methods for producing a `Vec<u8>` or `BytesMut` from
any type which impls `Encode`.
---
 ssh-encoding/src/encode.rs | 18 +++++++++++++++++-
 ssh-encoding/src/lib.rs    |  3 +++
 ssh-encoding/src/writer.rs | 11 +++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/ssh-encoding/src/encode.rs b/ssh-encoding/src/encode.rs
index 0879ad7c..b0d3f8bc 100644
--- a/ssh-encoding/src/encode.rs
+++ b/ssh-encoding/src/encode.rs
@@ -10,7 +10,7 @@ use core::str;
 use alloc::{string::String, vec::Vec};
 
 #[cfg(feature = "bytes")]
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 
 /// Encoding trait.
 ///
@@ -34,6 +34,22 @@ pub trait Encode {
         self.encoded_len()?.encode(writer)?;
         self.encode(writer)
     }
+
+    /// Encode this value, returning a `Vec<u8>` containing the encoded message.
+    #[cfg(feature = "alloc")]
+    fn encode_vec(&self) -> Result<Vec<u8>, Error> {
+        let mut ret = Vec::with_capacity(self.encoded_len()?);
+        self.encode(&mut ret)?;
+        Ok(ret)
+    }
+
+    /// Encode this value, returning a [`BytesMut`] containing the encoded message.
+    #[cfg(feature = "bytes")]
+    fn encode_bytes(&self) -> Result<BytesMut, Error> {
+        let mut ret = BytesMut::with_capacity(self.encoded_len()?);
+        self.encode(&mut ret)?;
+        Ok(ret)
+    }
 }
 
 /// Encode a single `byte` to the writer.
diff --git a/ssh-encoding/src/lib.rs b/ssh-encoding/src/lib.rs
index 74373949..3972909d 100644
--- a/ssh-encoding/src/lib.rs
+++ b/ssh-encoding/src/lib.rs
@@ -54,5 +54,8 @@ pub use crate::{
 #[cfg(feature = "base64")]
 pub use crate::{base64::Base64Reader, base64::Base64Writer};
 
+#[cfg(feature = "bytes")]
+pub use bytes;
+
 #[cfg(feature = "pem")]
 pub use crate::pem::{DecodePem, EncodePem};
diff --git a/ssh-encoding/src/writer.rs b/ssh-encoding/src/writer.rs
index 7b0fecab..41641da6 100644
--- a/ssh-encoding/src/writer.rs
+++ b/ssh-encoding/src/writer.rs
@@ -5,6 +5,9 @@ use crate::Result;
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
 
+#[cfg(feature = "bytes")]
+use bytes::{BufMut, BytesMut};
+
 #[cfg(feature = "sha2")]
 use sha2::{Digest, Sha256, Sha512};
 
@@ -23,6 +26,14 @@ impl Writer for Vec<u8> {
     }
 }
 
+#[cfg(feature = "bytes")]
+impl Writer for BytesMut {
+    fn write(&mut self, bytes: &[u8]) -> Result<()> {
+        self.put(bytes);
+        Ok(())
+    }
+}
+
 #[cfg(feature = "sha2")]
 impl Writer for Sha256 {
     fn write(&mut self, bytes: &[u8]) -> Result<()> {