From 527270812c5cfd0dcb93c48da6734aff7e0f548f Mon Sep 17 00:00:00 2001
From: Lars Erik Wik <lars.erik.wik@northern.tech>
Date: Mon, 21 Oct 2024 14:43:35 +0200
Subject: [PATCH] Fix potential bug where TLSRecv returns less than expected

TLSRecv may return less than the requested amount of Bytes. We fix this
my repeatedly calling it until we get the expected number.

Ticket: None
Changelog: None
Signed-off-by: Lars Erik Wik <lars.erik.wik@northern.tech>
---
 libcfnet/client_code.c | 18 ++++++++++++++----
 libcfnet/net.c         | 27 +++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/libcfnet/client_code.c b/libcfnet/client_code.c
index b888e6435d..0535c995fb 100644
--- a/libcfnet/client_code.c
+++ b/libcfnet/client_code.c
@@ -810,7 +810,7 @@ bool CopyRegularFileNet(const char *source, const char *dest, off_t size,
         assert(toget > 0);
 
         /* Stage C1 - receive */
-        int n_read;
+        int n_read = 0;
 
         const ProtocolVersion version = conn->conn_info->protocol;
 
@@ -820,7 +820,19 @@ bool CopyRegularFileNet(const char *source, const char *dest, off_t size,
         }
         else if (ProtocolIsTLS(version))
         {
-            n_read = TLSRecv(conn->conn_info->ssl, buf, toget);
+            /* TLSRecv may return less than the requested number of Bytes, in
+             * which case we repeatedly call it until we get the expected
+             * number. */
+            while (n_read < toget)
+            {
+                int rc = TLSRecv(conn->conn_info->ssl, buf + n_read, toget - n_read);
+                if (rc <= 0)
+                {
+                    n_read = rc;
+                    break;
+                }
+                n_read += rc;
+            }
         }
         else
         {
@@ -829,8 +841,6 @@ bool CopyRegularFileNet(const char *source, const char *dest, off_t size,
             n_read = -1;
         }
 
-        /* TODO what if 0 < n_read < toget? Might happen with TLS. */
-
         if (n_read <= 0)
         {
             /* This may happen on race conditions, where the file has shrunk
diff --git a/libcfnet/net.c b/libcfnet/net.c
index 9de7663f64..d0a0a0e5ba 100644
--- a/libcfnet/net.c
+++ b/libcfnet/net.c
@@ -158,7 +158,19 @@ int ReceiveTransaction(ConnectionInfo *conn_info, char *buffer, int *more)
         ret = RecvSocketStream(conn_info->sd, proto, CF_INBAND_OFFSET);
         break;
     case CF_PROTOCOL_TLS:
-        ret = TLSRecv(conn_info->ssl, proto, CF_INBAND_OFFSET);
+        /* TLSRecv may return less than the requested number of Bytes, in
+         * which case we repeatedly call it until we get the expected number.
+         */
+        while (ret < CF_INBAND_OFFSET)
+        {
+            int rc = TLSRecv(conn_info->ssl, proto + ret, CF_INBAND_OFFSET - ret);
+            if (rc <= 0)
+            {
+                ret = rc;
+                break;
+            }
+            ret += rc;
+        }
         break;
     default:
         UnexpectedError("ReceiveTransaction: ProtocolVersion %d!",
@@ -251,7 +263,18 @@ int ReceiveTransaction(ConnectionInfo *conn_info, char *buffer, int *more)
         ret = RecvSocketStream(conn_info->sd, buffer, len);
         break;
     case CF_PROTOCOL_TLS:
-        ret = TLSRecv(conn_info->ssl, buffer, len);
+        while (ret < len)
+        {
+            /* TLSRecv may return less than the requested number of Bytes, in
+             * which case we repeatedly call it until we get the expected
+             * number. */
+            int rc = TLSRecv(conn_info->ssl, buffer + ret, len - ret);
+            if (rc <= 0)
+            {
+                ret = rc;
+            }
+            ret += rc;
+        }
         break;
     default:
         UnexpectedError("ReceiveTransaction: ProtocolVersion %d!",