From 62f3c384b20688f481ac3f270e620734d540dd32 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Mon, 9 Aug 2021 17:21:22 +0100
Subject: [PATCH] Start `enc` socket wrapper
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for rsh's current commit: Curse − 凶
---
 src/cap.rs      | 13 ++++++-----
 src/sock.rs     | 43 ++++++++++++++++++++++++++++++++----
 src/sock/enc.rs | 58 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+), 9 deletions(-)
 create mode 100644 src/sock/enc.rs

diff --git a/src/cap.rs b/src/cap.rs
index a447080..016d40d 100644
--- a/src/cap.rs
+++ b/src/cap.rs
@@ -2,6 +2,7 @@
 use super::*;
 use std::time::Duration;
 
+/*
 pub mod fail;
 pub use fail::Failures;
 
@@ -18,23 +19,25 @@ pub enum Leniency
     /// Immediately disconnect the socket on **any** malformed/missed message.
     None,
 }
+ */
 
 /// A capability (permission) for a raw socket's data transmission.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum RawSockCapability
 {
-    /// Process
+    /// Process messages that aren't signed.
     AllowUnsignedMessages,
 
     /// Do not disconnect the socket when a malformed message is received, just ignore the message.
     SoftFail,
 
     /// Throttle the number of messages to process
-    RateLimit{ tx: usize, rx: usize },
+    RateLimit { tx: usize, rx: usize },
 
     /// The request response timeout for messages with an expected response.
     RecvRespTimeout { tx: Duration, rx: Duration },
+
+    /// Max number of bytes to read for a single message.
+    //TODO: Implement this for message
+    MaxMessageSize(usize),
 }
-/// A collection of `RawSockCapability`s for a specific connection.
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
-pub struct RawSockCapabilities;
diff --git a/src/sock.rs b/src/sock.rs
index 9296642..3237b8c 100644
--- a/src/sock.rs
+++ b/src/sock.rs
@@ -1,6 +1,6 @@
 //! Socket handlers
 use super::*;
-
+use std::collections::HashSet;
 use std::net::{
     SocketAddr,
 };
@@ -9,9 +9,16 @@ use tokio::io::{
     AsyncRead
 };
 use tokio::task::JoinHandle;
+use tokio::sync::{
+    mpsc,
+    oneshot,
+};
 use futures::Future;
+use bytes::Bytes;
+
 use cancel::*;
 
+pub mod enc;
 
 /// Details of a newly accepted raw socket peer.
 ///
@@ -32,20 +39,48 @@ pub struct RawSockPeerTrusted
     /// The socket's details
     pub sock_details: RawSockPeerAccepted,
 
-    ///
-    pub cap_allow: cap::RawSockCapabilities,
+    /// Capabilities for this peer
+    pub cap_allow: HashSet<cap::RawSockCapability>,
+}
+
+/// A raw, received message from a `RawPeer`.
+#[derive(Debug)]
+pub struct RawMessage{
+    message_bytes: Bytes,
+}
+
+/// A connected raw peer, created and handled by `handle_new_socket_with_shutdown()`.
+#[derive(Debug)]
+pub struct RawPeer{
+    info: RawSockPeerTrusted,
+    rx: mpsc::Receiver<RawMessage>,
 }
 
 /// Handles a **newly connected** raw  socket.
 ///
 /// This will handle setting up socket peer encryption and validation.
-pub fn handle_new_socket_with_shutdown<R, W, C: cancel::CancelFuture + 'static + Send>(sock_details: RawSockPeerAccepted, tx: W, rx: R, shutdown: C) -> JoinHandle<eyre::Result<()>>
+pub fn handle_new_socket_with_shutdown<R, W, C: cancel::CancelFuture + 'static + Send>(
+    sock_details: RawSockPeerAccepted,
+    set_peer: oneshot::Sender<RawPeer>,
+    tx: W, rx: R,
+    shutdown: C
+) -> JoinHandle<eyre::Result<()>>
 where R: AsyncRead + Unpin + Send + 'static,
       W: AsyncWrite + Unpin + Send + 'static
 {
     tokio::spawn(async move {
 	match {
 	    with_cancel!(async move {
+		// Create empty cap
+		let mut sock_details = RawSockPeerTrusted {
+		    sock_details,
+		    cap_allow: HashSet::new(),
+		};
+
+		// Set up encryption
+		
+		
+		//TODO: Find caps for this peer.
 		
 		Ok(())
 	    }, shutdown)
diff --git a/src/sock/enc.rs b/src/sock/enc.rs
new file mode 100644
index 0000000..fa57a40
--- /dev/null
+++ b/src/sock/enc.rs
@@ -0,0 +1,58 @@
+//! Socket encryption wrapper
+use super::*;
+use cryptohelpers::{
+    rsa::{
+	RsaPublicKey,
+	RsaPrivateKey,
+    },
+    sha256,
+};
+use chacha20stream::{
+    AsyncSink,
+};
+use std::sync::Arc;
+use tokio::{
+    sync::{
+	RwLock,
+    },
+    io::{
+	self,
+	DuplexStream,
+    },
+};
+
+/// Encrypted socket information.
+#[derive(Debug)]
+struct ESockInfo {
+    us: RsaPrivateKey,
+    them: Option<RsaPublicKey>,
+}
+
+/// A tx+rx socket.
+#[pin_project]
+#[derive(Debug)]
+pub struct ESock<W, R> {
+    info: RwLock<ESockInfo>,
+
+    #[pin]
+    // Raw (not encrypted) reader
+    rx: R,
+    #[pin]
+    tx: AsyncSink<W>,
+}
+
+/// Write half for `ESock`.
+#[pin_project]
+#[derive(Debug)]
+pub struct ESockWriteHalf<W>(Arc<ESockInfo>, #[pin] AsyncSink<W>);
+
+/// Read half for `ESock`.
+#[pin_project]
+#[derive(Debug)]
+pub struct ESockReadHalf<R>(
+    Arc<ESockInfo>,
+    
+    #[pin] R, // read from this (raw.)
+    #[pin] AsyncSink<DuplexStream>, // sink raw from `R`  here, outputs decrypted bytes into next.
+    #[pin] DuplexStream, // read decrypted bytes from here.
+);