From 9002b5688c4ad8a1db3a5ab23601c177610d45d8 Mon Sep 17 00:00:00 2001
From: kevinj <kevin.jerebica@gmail.com>
Date: Thu, 8 Aug 2024 13:25:23 +0200
Subject: [PATCH] fix: wait for socket transmissions to finish

---
 proxy.c   | 61 +++++++++++++++++++++++++++++++++++++------------------
 structs.h | 15 +++++++-------
 2 files changed, 48 insertions(+), 28 deletions(-)

diff --git a/proxy.c b/proxy.c
index b970808..f38815a 100644
--- a/proxy.c
+++ b/proxy.c
@@ -162,8 +162,13 @@ void do_err(void) {
 
 int do_fwd_clt(void) {
     int bytes = 0; 
+    int ret = 0; 
     do {
-        bytes += send(clt_sock, srv_msg+bytes, srv_msg_len-bytes, 0);
+        ret = send(clt_sock, srv_msg+bytes, srv_msg_len-bytes, 0);
+        if (ret < 0) 
+            return -1; 
+        
+        bytes += ret;
     } while (bytes < srv_msg_len);
 
     return 0;
@@ -176,20 +181,28 @@ int do_prs_srv(void) {
 }
 
 int do_rcv_srv(void) {
-    int ret = recv(srv_sock, srv_msg, PROXY_MAX_MSGLEN, 0);
-    if (ret < 0) 
-        return -1;
+    int bytes = 0;
+    int ret = 0; 
+    do {
+        ret = recv(srv_sock, srv_msg+bytes, PROXY_MAX_MSGLEN-bytes, 0);
+        if (ret < 0) 
+            return -1;
+        if (!ret) 
+            break;
 
-    srv_msg_len = ret;
+        bytes += ret;
+    } while (bytes < sizeof(PROXY_MAX_MSGLEN));
+
+    srv_msg_len = bytes;
 
     if (debug) {
-        fprintf(stdout, "[%d] Received server message: %s\n", statem, clt_msg);
+        fprintf(stdout, "[%d] Received server message: %s\n", statem, srv_msg);
     }
 
     return 0; 
 }
 
-int con_srv(void) {
+int do_con_srv(void) {
     int ret;
     char *host = getheader("Host");
     if (!host)
@@ -218,16 +231,14 @@ int con_srv(void) {
 }
 
 int do_fwd_srv(void) {
-    int ret = 0;
-    if (srv_sock < 0)
-        ret = con_srv();
-
-    if (ret < 0)
-        return -1;
-
     int bytes = 0;
+    int ret = 0;
     do {
-        bytes += send(srv_sock, clt_msg+bytes, clt_msg_len-bytes, 0);
+        ret = send(srv_sock, clt_msg+bytes, clt_msg_len-bytes, 0);
+        if (ret < 0) 
+            return -1; 
+
+        bytes += ret; 
     } while (bytes < clt_msg_len);
 
     return 0;
@@ -255,12 +266,19 @@ int do_prs_clt(void) {
 }
 
 int do_rcv_clt(void) {
-    int ret; 
-    ret = recv(clt_sock, clt_msg, PROXY_MAX_MSGLEN, 0);
-    if (ret < 0)
-        return -1;
+    int bytes = 0; 
+    int ret = 0;
+    do {
+        ret = recv(clt_sock, clt_msg+bytes, PROXY_MAX_MSGLEN-bytes, 0);
+        if (ret < 0) 
+            return -1;
+        if (!ret) 
+            break;
 
-    clt_msg_len = ret;
+        bytes += ret;
+    } while (bytes < sizeof(PROXY_MAX_MSGLEN));
+
+    clt_msg_len = bytes;
 
     if (debug) {
         fprintf(stdout, "[%d] Received client message: %s\n", statem, clt_msg);
@@ -309,6 +327,9 @@ void dostatem() {
         case STATEM_PRS_CLT: 
             ret = do_prs_clt();
             break; 
+        case STATEM_CON_SRV: 
+            ret = do_con_srv();
+            break;
         case STATEM_FWD_SRV: 
             ret = do_fwd_srv();
             break; 
diff --git a/structs.h b/structs.h
index 6e23440..ce2b0b6 100644
--- a/structs.h
+++ b/structs.h
@@ -15,17 +15,16 @@
                     "[ ]+([a-zA-Z0-9_.,/-]+)[\n\r]*$"
 #define REGEX_HEADER "^([a-zA-Z0-9_-]*):[ \t]+([^\r\n]*)"
 
-#define CLIENT_MESSAGE "GET http://google.com/auth HTTP/1.0\n"\
-                       "\n"\
-                       "Host: google.com\n"\
-                       "Authorization: Bearer ffja2439gjawgjgojserg\n"
+#define CLIENT_MESSAGE "GET http://archive.0xdeadbeer.org/ HTTP/1.0\r\n\r\n"\
+                       "Host: archive.0xdeadbeer.org\r\n"\
 
 #define STATEM_RCV_CLT 0b00000001
 #define STATEM_PRS_CLT 0b00000010
-#define STATEM_FWD_SRV 0b00000100
-#define STATEM_RCV_SRV 0b00001000
-#define STATEM_PRS_SRV 0b00010000
-#define STATEM_FWD_CLT 0b00100000
+#define STATEM_CON_SRV 0b00000100
+#define STATEM_FWD_SRV 0b00001000
+#define STATEM_RCV_SRV 0b00010000
+#define STATEM_PRS_SRV 0b00100000
+#define STATEM_FWD_CLT 0b01000000
 #define STATEM_ERR     0b10000000
 
 struct header {