From 907862f5540553e7a28a5445683498ffb30a86d3 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Fri, 13 Aug 2021 16:08:37 +0100
Subject: [PATCH] Source: Encrypt and Decrypt tests passing. (For both internal
 and external buffer.)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for chacha20stream's current commit: Half blessing − 半吉
---
 src/stream/mod.rs    | 63 ++++++++++++++++++++++++++++++++++++++++++--
 src/stream/sink.rs   |  2 --
 src/stream/source.rs | 40 ++++++++++++++++++++++++++--
 3 files changed, 99 insertions(+), 6 deletions(-)

diff --git a/src/stream/mod.rs b/src/stream/mod.rs
index f1e9073..c6eee3d 100644
--- a/src/stream/mod.rs
+++ b/src/stream/mod.rs
@@ -31,6 +31,7 @@ pub use source::Source;
 mod tests
 {
     use super::*;
+    use std::io::Cursor;
 
     const INPUT: &'static str = "Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!Hello world!";
 
@@ -50,6 +51,20 @@ mod tests
 	stream
     }
 
+    /// Create a source from a slice of bytes with this key and IV.
+    fn create_source<'a, T: ?Sized + AsRef<[u8]> +'a>(source: &'a T, key: Key, iv: IV, enc: bool) -> Source<Cursor<&'a [u8]>>
+    {
+	eprintln!("({}) Key: {}, IV: {}, Input: ({}: {})", ["dec", "enc"][enc as usize], key, iv, source.as_ref().len(), source.as_ref().hex());
+	
+	let stream = if enc {
+	    Source::encrypt(Cursor::new(source.as_ref()), key, iv)
+	} else {
+	    Source::decrypt(Cursor::new(source.as_ref()), key, iv)
+	}.expect("sink::enc");
+	
+	stream
+    }
+
     #[test]
     fn enc()
     {
@@ -58,6 +73,50 @@ mod tests
 	eprintln!("Sink ends: {:?}", enc_stream(INPUT.as_bytes(), key, iv));
     }
 
+    #[test]
+    fn source_enc()
+    {
+	let (key, iv) = cha::keygen();
+
+	const INPUT: &'static [u8] = b"Hello world!";
+	println!("Input ({} bytes, hex): {}", INPUT.len(), INPUT.hex());
+
+	let mut source = create_source(INPUT, key, iv, true);
+
+	let mut output = Vec::with_capacity(INPUT.len());
+	io::copy(&mut source, &mut output).expect("Failed to copy source to output");
+
+	println!("Output ({} bytes, hex): {}", output.len(), output.hex());
+    }
+
+
+    #[test]
+    fn source_dec()
+    {
+	let (key, iv) = cha::keygen();
+
+	const INPUT: &'static [u8] = b"Hello world!";
+	println!("Input ({} bytes, hex): {}", INPUT.len(), INPUT.hex());
+
+	let mut source = create_source(INPUT, key, iv, true);
+
+	let mut temp = Vec::with_capacity(INPUT.len());
+	io::copy(&mut source, &mut temp).expect("Failed to copy source to output (encrypt)");
+
+	println!("Encrypted ({} bytes, hex): {}", temp.len(), temp.hex());
+
+	// decrypt
+
+	let mut source = create_source(&mut temp, key, iv, false);
+
+	let mut temp = Vec::with_capacity(INPUT.len());
+	io::copy(&mut source, &mut temp).expect("Failed to copy source to output (decrypt)");
+	
+	println!("Decrypted ({} bytes, hex): {}", temp.len(), temp.hex());
+
+	assert_eq!(INPUT, &temp[..]);
+    }
+    
     #[test]
     fn dec()
     {
@@ -78,9 +137,9 @@ mod tests
 	}
 	assert_eq!(&dec_buffer[..], INPUT.as_bytes());
     }
-
+    
     /// Checks if explicit clear is actually clearing.
-    #[cfg(feature="explicit_clear")] 
+	#[cfg(feature="explicit_clear")] 
     #[test]
     fn remainder()
     {
diff --git a/src/stream/sink.rs b/src/stream/sink.rs
index dc4fef6..baa124e 100644
--- a/src/stream/sink.rs
+++ b/src/stream/sink.rs
@@ -144,8 +144,6 @@ where W: Write
     {
 	(self.stream, self.crypter)
     }
-    
-
 }
 
 
diff --git a/src/stream/source.rs b/src/stream/source.rs
index 1412c55..37c3641 100644
--- a/src/stream/source.rs
+++ b/src/stream/source.rs
@@ -206,10 +206,46 @@ where R: Read
 
 }
 
+impl<R, K: ?Sized + BufferKind> Source<R, K>
+where R: Read
+{
+    /// Create a new Chacha Source stream wrapper from a reader
+    #[inline] fn new(stream: R, crypter: Crypter) -> Self
+    {
+	Self{stream, crypter, buffer: K::create_buffer(0)}
+    }
+
+    /// Create an encrypting Chacha Source stream wrapper
+    pub fn encrypt(stream: R, key: Key, iv: IV) -> Result<Self, Error>
+    {
+	Ok(Self::new(stream, cha::encrypter(key, iv)?))
+    }
+    
+    /// Create a decrypting Chacha Source stream wrapper
+    pub fn decrypt(stream: R, key: Key, iv: IV) -> Result<Self, Error>
+    {
+	Ok(Self::new(stream, cha::decrypter(key, iv)?))
+    }
+    
+
+    /// Consume into the inner stream
+    #[inline] pub fn into_inner(self) -> R
+    {
+	self.stream
+    }
+
+    /// Consume into the inner stream and crypter
+    #[inline] pub fn into_parts(self) -> (R, Crypter)
+    {
+	(self.stream, self.crypter)
+    }
+}
+
+
 impl<R> Source<R, UseBufferExternal>
 {
-    /// Convert this instance to use external buffer (instead of internal.)
-    pub fn with_reused_buffer(self) -> Source<R, UseBufferInternal>
+    /// Convert this instance to use internal buffer (instead of external.)
+    pub fn with_internal_buffer(self) -> Source<R, UseBufferInternal>
     {
 	Source {
 	    buffer: UseBufferInternal::create_buffer(UseBufferExternal::buffer_cap(&self)),