From 667ca0c035d167a34cfd6ea2a59017e8d2f8f33f Mon Sep 17 00:00:00 2001
From: Isabelle L <me@izzabelle.dev>
Date: Thu, 14 May 2020 00:23:41 -0500
Subject: [PATCH] added `read_encrypted` and `write_encrypted` functions

---
 src/lib.rs | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 308eedf..ccff7d3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@
 #![allow(dead_code)]
 
 use futures_util::io::{AsyncReadExt, AsyncWriteExt};
+use orion::aead;
 use std::convert::TryInto;
 use std::marker::Unpin;
 
@@ -49,6 +50,8 @@ impl Packet {
 }
 
 /// reads a `Packet` from a stream
+///
+/// if `Ok(None)` is returned the stream has been disconnected.
 pub async fn read<S>(stream: &mut S) -> Result<Option<Packet>>
 where
     S: AsyncReadExt + Unpin,
@@ -70,6 +73,23 @@ where
     Ok(Some(packet))
 }
 
+/// reads a `Packet` from a stream and decrypts
+///
+/// if `Ok(None)` is returned the stream has been disconnected.
+pub async fn read_encrypted<S>(stream: &mut S, key: &aead::SecretKey) -> Result<Option<Packet>>
+where
+    S: AsyncReadExt + Unpin,
+{
+    let packet = read(stream).await?;
+    match packet {
+        None => Ok(packet),
+        Some(mut packet) => {
+            packet.contents = aead::open(&key, &packet.contents)?;
+            Ok(Some(packet))
+        }
+    }
+}
+
 /// Writes a `Sendable` packet to a stream
 pub async fn write<S, P>(stream: &mut S, packet: P) -> Result<()>
 where
@@ -81,6 +101,19 @@ where
     Ok(())
 }
 
+/// Writes an encrypted `Sendable` packet to a stream
+pub async fn write_encrypted<S, P>(stream: &mut S, packet: P, key: &aead::SecretKey) -> Result<()>
+where
+    S: AsyncWriteExt + Unpin,
+    P: Sendable,
+{
+    let mut packet = packet.to_packet()?;
+    packet.contents = aead::seal(&key, &packet.contents)?;
+    let network_packet = packet.to_network_packet();
+    stream.write(&network_packet.0).await?;
+    Ok(())
+}
+
 /// Kinds of packets that can be sent
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 #[repr(u8)]