From 9ae9ff30cbefd84f20c9dc47ac476d32b31a717d Mon Sep 17 00:00:00 2001
From: Igor Klevanets <8238332+cerevra@users.noreply.github.com>
Date: Wed, 23 Sep 2020 23:15:01 +0300
Subject: [PATCH] agent.c: Recv and send all bytes via network in
 agent_transact_unix() (#510)

Files: agent.c

Notes:
Handle sending/receiving partial packet replies in agent.c API.

Credit: Klevanets Igor <cerevra@yandex-team.ru>
---
 src/agent.c | 48 ++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/agent.c b/src/agent.c
index c9efe304..7ce39707 100644
--- a/src/agent.c
+++ b/src/agent.c
@@ -175,6 +175,40 @@ agent_connect_unix(LIBSSH2_AGENT *agent)
     return LIBSSH2_ERROR_NONE;
 }
 
+#define RECV_SEND_ALL(func, socket, buffer, length, flags, abstract) \
+    int rc; \
+    size_t finished; \
+    \
+    finished = 0; \
+    \
+    while(finished < length) { \
+        rc = func(socket, \
+                  (char *)buffer + finished, length - finished, \
+                  flags, abstract); \
+        if(rc < 0) \
+            return rc; \
+        \
+        finished += rc; \
+    } \
+    \
+    return finished;
+
+static ssize_t _send_all(LIBSSH2_SEND_FUNC(func), libssh2_socket_t socket,
+                         const void *buffer, size_t length,
+                         int flags, void **abstract)
+{
+    RECV_SEND_ALL(func, socket, buffer, length, flags, abstract);
+}
+
+static ssize_t _recv_all(LIBSSH2_RECV_FUNC(func), libssh2_socket_t socket,
+                         void *buffer, size_t length,
+                         int flags, void **abstract)
+{
+    RECV_SEND_ALL(func, socket, buffer, length, flags, abstract);
+}
+
+#undef RECV_SEND_ALL
+
 static int
 agent_transact_unix(LIBSSH2_AGENT *agent, agent_transaction_ctx_t transctx)
 {
@@ -184,7 +218,8 @@ agent_transact_unix(LIBSSH2_AGENT *agent, agent_transaction_ctx_t transctx)
     /* Send the length of the request */
     if(transctx->state == agent_NB_state_request_created) {
         _libssh2_htonu32(buf, transctx->request_len);
-        rc = LIBSSH2_SEND_FD(agent->session, agent->fd, buf, sizeof buf, 0);
+        rc = _send_all(agent->session->send, agent->fd,
+                       buf, sizeof buf, 0, &agent->session->abstract);
         if(rc == -EAGAIN)
             return LIBSSH2_ERROR_EAGAIN;
         else if(rc < 0)
@@ -195,8 +230,8 @@ agent_transact_unix(LIBSSH2_AGENT *agent, agent_transaction_ctx_t transctx)
 
     /* Send the request body */
     if(transctx->state == agent_NB_state_request_length_sent) {
-        rc = LIBSSH2_SEND_FD(agent->session, agent->fd, transctx->request,
-                           transctx->request_len, 0);
+        rc = _send_all(agent->session->send, agent->fd, transctx->request,
+                       transctx->request_len, 0, &agent->session->abstract);
         if(rc == -EAGAIN)
             return LIBSSH2_ERROR_EAGAIN;
         else if(rc < 0)
@@ -207,7 +242,8 @@ agent_transact_unix(LIBSSH2_AGENT *agent, agent_transaction_ctx_t transctx)
 
     /* Receive the length of a response */
     if(transctx->state == agent_NB_state_request_sent) {
-        rc = LIBSSH2_RECV_FD(agent->session, agent->fd, buf, sizeof buf, 0);
+        rc = _recv_all(agent->session->recv, agent->fd,
+                       buf, sizeof buf, 0, &agent->session->abstract);
         if(rc < 0) {
             if(rc == -EAGAIN)
                 return LIBSSH2_ERROR_EAGAIN;
@@ -225,8 +261,8 @@ agent_transact_unix(LIBSSH2_AGENT *agent, agent_transaction_ctx_t transctx)
 
     /* Receive the response body */
     if(transctx->state == agent_NB_state_response_length_received) {
-        rc = LIBSSH2_RECV_FD(agent->session, agent->fd, transctx->response,
-                           transctx->response_len, 0);
+        rc = _recv_all(agent->session->recv, agent->fd, transctx->response,
+                       transctx->response_len, 0, &agent->session->abstract);
         if(rc < 0) {
             if(rc == -EAGAIN)
                 return LIBSSH2_ERROR_EAGAIN;