From c8a7a67ad32777040ba102607c18f66e118f91a5 Mon Sep 17 00:00:00 2001
From: David Venhoek <david@tweedegolf.com>
Date: Fri, 8 Mar 2024 12:27:59 +0100
Subject: [PATCH] Fix incorrect shutdown of tls streams.

---
 statime-linux/src/ke/client.rs | 8 +++++++-
 statime-linux/src/ke/server.rs | 5 +++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/statime-linux/src/ke/client.rs b/statime-linux/src/ke/client.rs
index b9c4ec349..6ef870ea7 100644
--- a/statime-linux/src/ke/client.rs
+++ b/statime-linux/src/ke/client.rs
@@ -1,7 +1,10 @@
 use std::{error::Error, io, sync::Arc};
 
 use rustls::{pki_types::ServerName, ClientConfig};
-use tokio::{io::AsyncReadExt, net::TcpStream};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
+};
 use tokio_rustls::TlsConnector;
 
 use super::record::{AssociationMode, NextProtocols, PtpKeyRequestMessage, PtpKeyResponseMessage};
@@ -27,6 +30,7 @@ pub async fn fetch_data(
     let mut stream = connector.connect(dnsname, stream).await?;
 
     request.write(&mut stream).await?;
+    stream.flush().await?;
 
     // we expect the to receive messages to be smaller than data_buf
     let mut data_buf = vec![0; 4096];
@@ -47,6 +51,8 @@ pub async fn fetch_data(
         }
     };
 
+    stream.shutdown().await?;
+
     let records: Vec<_> = records.into_iter().map(|r| r.into_owned()).collect();
 
     Ok(records.try_into()?)
diff --git a/statime-linux/src/ke/server.rs b/statime-linux/src/ke/server.rs
index 133e56a6f..1f271aaac 100644
--- a/statime-linux/src/ke/server.rs
+++ b/statime-linux/src/ke/server.rs
@@ -9,7 +9,7 @@ use std::{
 use log::{debug, info, warn, LevelFilter};
 use rustls::ServerConfig;
 use tokio::{
-    io::AsyncReadExt,
+    io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
     sync::RwLock,
     time::Instant,
@@ -214,7 +214,8 @@ async fn handle_connection(
 
     let keyset = store.read().await;
     let resp = respond(records, &keyset, ke_config).await?;
-    resp.write(stream).await?;
+    resp.write(&mut stream).await?;
+    stream.shutdown().await?;
     Ok(())
 }