From 2c0883878d22ee6d223172201fe227619df47e89 Mon Sep 17 00:00:00 2001
From: Chip Senkbeil <chip@senkbeil.org>
Date: Thu, 11 Nov 2021 13:03:36 -0600
Subject: [PATCH] Add --key-from-stdin option to listen cli command to read key
 from stdin instead of generating

---
 distant-core/src/net/mod.rs |  6 ++++++
 src/opt.rs                  |  6 ++++++
 src/subcommand/listen.rs    | 15 +++++++++++++--
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/distant-core/src/net/mod.rs b/distant-core/src/net/mod.rs
index 8ac958e..47d0556 100644
--- a/distant-core/src/net/mod.rs
+++ b/distant-core/src/net/mod.rs
@@ -81,6 +81,12 @@ impl<const N: usize> SecretKey<N> {
     }
 }
 
+impl<const N: usize> From<[u8; N]> for SecretKey<N> {
+    fn from(arr: [u8; N]) -> Self {
+        Self(arr)
+    }
+}
+
 impl<const N: usize> FromStr for SecretKey<N> {
     type Err = SecretKeyError;
 
diff --git a/src/opt.rs b/src/opt.rs
index 5e27487..d5e1de1 100644
--- a/src/opt.rs
+++ b/src/opt.rs
@@ -580,6 +580,12 @@ pub struct ListenSubcommand {
     #[structopt(long, default_value = &SERVER_CONN_MSG_CAPACITY_STR)]
     pub max_msg_capacity: u16,
 
+    /// If specified, the server will not generate a key but instead listen on stdin for the next
+    /// 32 bytes that it will use as the key instead. Receiving less than 32 bytes before stdin
+    /// is closed is considered an error and any bytes after the first 32 are not used for the key
+    #[structopt(long)]
+    pub key_from_stdin: bool,
+
     /// The time in seconds before shutting down the server if there are no active
     /// connections. The countdown begins once all connections have closed and
     /// stops when a new connection is made. In not specified, the server will not
diff --git a/src/subcommand/listen.rs b/src/subcommand/listen.rs
index 21c789e..7d373b7 100644
--- a/src/subcommand/listen.rs
+++ b/src/subcommand/listen.rs
@@ -8,12 +8,13 @@ use distant_core::{
 };
 use log::*;
 use tokio::{
-    io::{self, AsyncWriteExt},
+    io::{self, AsyncReadExt, AsyncWriteExt},
     task::JoinError,
 };
 
 #[derive(Debug, Display, Error, From)]
 pub enum Error {
+    BadKey,
     ConverToIpAddr(ConvertToIpAddrError),
     Fork,
     Io(io::Error),
@@ -23,6 +24,7 @@ pub enum Error {
 impl ExitCodeError for Error {
     fn to_exit_code(&self) -> ExitCode {
         match self {
+            Self::BadKey => ExitCode::Usage,
             Self::ConverToIpAddr(_) => ExitCode::NoHost,
             Self::Fork => ExitCode::OsErr,
             Self::Io(x) => x.to_exit_code(),
@@ -99,7 +101,16 @@ async fn run_async(cmd: ListenSubcommand, _opt: CommonOpt, is_forked: bool) -> R
     }
 
     // Bind & start our server
-    let key = SecretKey32::default();
+    let key = if cmd.key_from_stdin {
+        let mut buf = [0u8; 32];
+        let n = io::stdin().read_exact(&mut buf).await?;
+        if n < buf.len() {
+            return Err(Error::BadKey);
+        }
+        SecretKey32::from(buf)
+    } else {
+        SecretKey32::default()
+    };
     let key_hex_string = key.unprotected_to_hex_key();
     let codec = XChaCha20Poly1305Codec::from(key);