From bba6c8f39b92d2701b6622c2722914b2c8ac47bd Mon Sep 17 00:00:00 2001
From: Martin Hoffmann <martin@nlnetlabs.nl>
Date: Fri, 21 Jan 2022 17:46:24 +0100
Subject: [PATCH] RTR: Flush socket after final write before waiting.

This is necessary for TLS as it otherwise waits for more data to fill
its buffers.

This fixes a hanging bug in the RTR-over-TLS server.
---
 src/rtr/client.rs | 4 +++-
 src/rtr/server.rs | 7 +++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/rtr/client.rs b/src/rtr/client.rs
index 68e3f414..877b81bc 100644
--- a/src/rtr/client.rs
+++ b/src/rtr/client.rs
@@ -12,7 +12,7 @@ use std::{error, fmt, io};
 use std::future::Future;
 use std::marker::Unpin;
 use tokio::time::{timeout, timeout_at, Duration, Instant};
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use super::payload::{Action, Payload, Timing};
 use super::pdu;
 use super::state::State;
@@ -273,6 +273,7 @@ where
         pdu::SerialQuery::new(
             self.version(), state,
         ).write(&mut self.sock).await?;
+        self.sock.flush().await?;
         let start = match self.try_io(FirstReply::read).await? {
             FirstReply::Response(start) => start,
             FirstReply::Reset(_) => {
@@ -324,6 +325,7 @@ where
         pdu::ResetQuery::new(
             self.version()
         ).write(&mut self.sock).await?;
+        self.sock.flush().await?;
         let start = self.try_io(|sock| {
             pdu::CacheResponse::read(sock)
         }).await?;
diff --git a/src/rtr/server.rs b/src/rtr/server.rs
index bd2d90c8..0b84bfff 100644
--- a/src/rtr/server.rs
+++ b/src/rtr/server.rs
@@ -11,7 +11,7 @@ use futures_util::future;
 use futures_util::pin_mut;
 use futures_util::future::Either;
 use log::debug;
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::sync::broadcast;
 use tokio::task::spawn;
 use tokio_stream::{Stream, StreamExt};
@@ -401,6 +401,7 @@ impl<Sock: Socket, Source: PayloadSource> Connection<Sock, Source> {
                 pdu::EndOfData::new(
                     self.version(), state, timing
                 ).write(&mut self.sock).await?;
+                self.sock.flush().await?;
                 self.sock.update(state, true);
                 Ok(())
             }
@@ -436,6 +437,7 @@ impl<Sock: Socket, Source: PayloadSource> Connection<Sock, Source> {
         pdu::EndOfData::new(
             self.version(), state, timing
         ).write(&mut self.sock).await?;
+        self.sock.flush().await?;
         self.sock.update(state, true);
         Ok(())
     }
@@ -444,7 +446,8 @@ impl<Sock: Socket, Source: PayloadSource> Connection<Sock, Source> {
     async fn error(
         &mut self, err: pdu::Error
     ) -> Result<(), io::Error> {
-        err.write(&mut self.sock).await
+        err.write(&mut self.sock).await?;
+        self.sock.flush().await
     }
 
     /// Sends a serial notify query.