From e539ab7a082854892c8f7722d2ec5b6b14d3cf24 Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Sun, 12 Jul 2020 21:32:43 +0200 Subject: ssh: Thus, look around sshnc_maybe_recv() --- ssh/client_proxy.c | 31 ++++++++++++++++++++----------- ssh/ssh_client.c | 28 ++++++++++++++++++---------- 2 files changed, 38 insertions(+), 21 deletions(-) diff --git a/ssh/client_proxy.c b/ssh/client_proxy.c index 58f40cb..1ce5e12 100644 --- a/ssh/client_proxy.c +++ b/ssh/client_proxy.c @@ -88,20 +88,29 @@ static void* proxy_thread_entry(void *thread_data_) { } if (polls[0].revents & (POLLIN | POLLHUP)) { - char buffer[4096]; - size_t length = 0; - ret = sshnc_maybe_recv(client, sizeof buffer, buffer, &length); - if (ret == SSHNC_OK) { - if (writeall(client_sock, buffer, length) < 0) { - // Error writing back to downstream, let's just close + // Get all data currently available + bool should_exit = false; + while (true) { + char buffer[4096]; + size_t length = 0; + ret = sshnc_maybe_recv(client, sizeof buffer, buffer, &length); + if (ret == SSHNC_AGAIN) break; + if (ret == SSHNC_OK) { + if (writeall(client_sock, buffer, length) < 0) { + // Error writing back to downstream, let's just close + should_exit = true; + break; + } + } else if (ret == SSHNC_EOF) { + should_exit = true; break; + } else { + fprintf(stderr, "Error on SSH recv: %s\n", sshnc_strerror(ret)); + goto cleanup; } - } else if (ret == SSHNC_EOF) { - break; - } else if (ret != SSHNC_AGAIN) { - fprintf(stderr, "Error on SSH recv: %s\n", sshnc_strerror(ret)); - goto cleanup; } + + if (should_exit) break; } if (polls[1].revents & (POLLIN | POLLHUP)) { diff --git a/ssh/ssh_client.c b/ssh/ssh_client.c index 5c7f084..3e4a7c6 100644 --- a/ssh/ssh_client.c +++ b/ssh/ssh_client.c @@ -110,17 +110,25 @@ int main(int argc, char **argv) { } if (polls[0].revents & (POLLIN | POLLHUP)) { - char buffer[4096]; - size_t length = 0; - ret = sshnc_maybe_recv(client, sizeof buffer, buffer, &length); - if (ret == SSHNC_OK) { - fwrite(buffer, 1, length, stdout); - } else if (ret == SSHNC_EOF) { - break; - } else if (ret != SSHNC_AGAIN) { - fprintf(stderr, "Error on SSH recv: %s\n", sshnc_strerror(ret)); - goto cleanup; + // Get all data currently available + bool should_exit = false; + while (true) { + char buffer[4096]; + size_t length = 0; + ret = sshnc_maybe_recv(client, sizeof buffer, buffer, &length); + if (ret == SSHNC_AGAIN) break; + if (ret == SSHNC_EOF) { + should_exit = true; + break; + } else if (ret == SSHNC_OK) { + fwrite(buffer, 1, length, stdout); + } else { + fprintf(stderr, "Error on SSH recv: %s\n", sshnc_strerror(ret)); + goto cleanup; + } } + + if (should_exit) break; } if (polls[1].revents & (POLLIN | POLLHUP)) { -- cgit v1.2.3-70-g09d2