From 8a1407185a88825fdab905d97592a69cacee2432 Mon Sep 17 00:00:00 2001
From: "Christoph J. Scherr" <Christoph.Scherr@newtec.de>
Date: Tue, 23 Jan 2024 18:21:14 +0100
Subject: [PATCH] a little tls error handling

---
 src/server/errors.rs |  8 ++++----
 src/server/mod.rs    | 42 +++++++++++++++++++++++++++++++-----------
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/src/server/errors.rs b/src/server/errors.rs
index 90dd1f5..0b661c9 100644
--- a/src/server/errors.rs
+++ b/src/server/errors.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, string::FromUtf8Error};
+use std::{fmt::Display, str::Utf8Error, string::FromUtf8Error};
 
 use anyhow;
 use thiserror::Error;
@@ -11,7 +11,7 @@ pub enum ServerError {
     Timeout(Elapsed),
     Anyhow(anyhow::Error),
     IO(std::io::Error),
-    Format(FromUtf8Error),
+    Format(Utf8Error),
 }
 
 impl From<anyhow::Error> for ServerError {
@@ -26,8 +26,8 @@ impl From<std::io::Error> for ServerError {
     }
 }
 
-impl From<FromUtf8Error> for ServerError {
-    fn from(value: FromUtf8Error) -> Self {
+impl From<Utf8Error> for ServerError {
+    fn from(value: Utf8Error) -> Self {
         Self::Format(value)
     }
 }
diff --git a/src/server/mod.rs b/src/server/mod.rs
index aad6334..f91a87e 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -1,13 +1,18 @@
 #![cfg(feature = "server")]
 use std::{
-    fs::File, net::SocketAddr, sync::{atomic::AtomicUsize, Arc}, time::Duration
+    fs::File,
+    net::SocketAddr,
+    sync::{atomic::AtomicUsize, Arc},
+    time::Duration,
 };
 
-use libpt::log::{debug, info, trace, warn};
+use libpt::log::{debug, error, info, trace, warn};
 use rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use rustls_pemfile::{certs, private_key};
 use tokio::{
-    io::{split, AsyncReadExt, AsyncWriteExt, BufReader}, net::{TcpListener, TcpStream}, time::{self, timeout}
+    io::{split, AsyncReadExt, AsyncWriteExt, BufReader},
+    net::{TcpListener, TcpStream},
+    time::{self, timeout},
 };
 use tokio_rustls::{rustls, TlsAcceptor};
 
@@ -29,7 +34,7 @@ pub struct Server {
 impl Server {
     pub async fn build(cfg: Config) -> anyhow::Result<Self> {
         let certs = Self::load_certs(cfg.clone())?;
-        let key = Self::load_key(cfg.clone())?.unwrap();
+        let key = Self::load_key(cfg.clone())?.expect("bad key?");
         let tls_config = rustls::ServerConfig::builder()
             .with_no_client_auth()
             .with_single_cert(certs, key)?;
@@ -69,12 +74,18 @@ impl Server {
                 }
             };
             let ref_self = rc_self.clone();
-            let acceptor = rc_self.acceptor.clone();
             // NOTE: we can only start the task now. If we start it before accepting connections
             // (so that the task theoretically accepts the connection), we would create endless
             // tasks in a loop.
             tokio::spawn(async move {
-                let stream: tokio_rustls::server::TlsStream<_> = acceptor.accept(stream).await.unwrap();
+                let stream: tokio_rustls::server::TlsStream<_> =
+                    match ref_self.acceptor.accept(stream).await {
+                        Ok(s) => s,
+                        Err(err) => {
+                            error!("could not accept tcp stream: {err}");
+                            return;
+                        }
+                    };
                 ref_self.peer_add(1);
                 match ref_self.handle_stream(stream, addr).await {
                     Ok(_) => (),
@@ -116,19 +127,25 @@ impl Server {
         )
     }
 
-    async fn handle_stream(&self, stream: tokio_rustls::server::TlsStream<TcpStream>, addr: SocketAddr) -> Result<()> {
-        let mut pings: usize = 0;
+    async fn handle_stream(
+        &self,
+        stream: tokio_rustls::server::TlsStream<TcpStream>,
+        addr: SocketAddr,
+    ) -> Result<()> {
         debug!("new peer: {:?}", addr);
         let mut buf = [0; BUF_SIZE];
         let (mut reader, mut writer) = split(stream);
         loop {
             match reader.read(&mut buf).await {
-                Ok(len) if len == 0 => { break;},
+                Ok(len) if len == 0 => {
+                    break;
+                }
                 Ok(_) => (),
                 Err(err) => {
                     eprintln!("reader.read err: {err}")
                 }
             }
+            debug!("< {addr:?} : \"{}\"", self.decode(&buf)?);
 
             writer.write(b"pong\0").await?;
 
@@ -140,7 +157,10 @@ impl Server {
     }
 
     #[inline]
-    fn decode(&self, buf: &Vec<u8>) -> Result<String> {
-        Ok(String::from_utf8(buf.clone())?.replace('\n', "\\n"))
+    fn decode(&self, buf: &[u8]) -> Result<String> {
+        match std::str::from_utf8(buf) {
+            Ok(s) => Ok(s.to_string()),
+            Err(err) => Err(err.into()),
+        }
     }
 }