ssh: Thus, look around sshnc_maybe_recv()

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)) {