Fix incorrect shutdown of tls streams.

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()?)

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(())
 }