fix buffer overflow vulnerabilityHEADmaster

1 files changed, 28 insertions, 15 deletions

diff --git a/src/main.c b/src/main.c
index 8817b83..c180c7e 100644
--- a/src/main.c
+++ b/src/main.c
@@ -8,27 +8,28 @@
 #include "log.h"
 #include "server.h"
 
-#define BUF_CAP 80000
+#define BUF_CAP 65536
 #define COMMON_CAP 1024
+#define REQ_PATH_CAP 512
+
 #define PORT "8079"
 
-#define RM_LF(str) do {                                     \
-        signed long len = strlen(str)-1;                    \
-        if(len >= 0 && str[len] == '\n') str[len] = '\0';   \
-    } while(0)
+#define MIN(a,b) (((a)<(b))?(a):(b))
 
-#define _SEND_BUF_ADD(str) do {                             \
-        memcpy(&send_buf[send_buf_sz], str, strlen(str));   \
-        send_buf_sz += strlen(str);                         \
+#define _SEND_BUF_ADD(str) do {                                      \
+        size_t n = MIN(strlen(str), BUF_CAP - (size_t)send_buf_sz);  \
+        memcpy(&send_buf[send_buf_sz], str, n);                      \
+        send_buf_sz += n;                                            \
     } while(0)
 #define _SEND_BUF_ADD_LINE(str) do {                  \
         _SEND_BUF_ADD(str);                           \
         _SEND_BUF_ADD("\r\n");                        \
     } while(0)
 
-#define SEND_BUF_ADD(str) do {                               \
-        memcpy(&send_buf[*send_buf_sz], str, strlen(str));   \
-        *send_buf_sz += strlen(str);                         \
+#define SEND_BUF_ADD(str) do {                                       \
+        size_t n = MIN(strlen(str), BUF_CAP - (size_t)*send_buf_sz); \
+        memcpy(&send_buf[*send_buf_sz], str, strlen(str));           \
+        *send_buf_sz += n;                                           \
     } while(0)
 #define SEND_BUF_ADD_LINE(str) do {                  \
         SEND_BUF_ADD(str);                           \
@@ -37,7 +38,7 @@
 
 #define SEND_BUF_ADD_LINE_LONG(lon) do {        \
         char str[COMMON_CAP];                   \
-        sprintf(str, "%ld", lon);               \
+        snprintf(str, COMMON_CAP, "%ld", lon);  \
         SEND_BUF_ADD_LINE(str);                 \
     } while(0)
 
@@ -64,6 +65,11 @@ static int get_content_type(char *file_path, char *content_type)
         goto close;
     }
 
+    if(strlen(mime) >= COMMON_CAP) {
+        err("get_content_type: mime info too long");
+        goto close;
+    }
+
     memcpy(content_type, mime, strlen(mime));
 
     ret = 0;
@@ -82,6 +88,13 @@ static int get_content_type(char *file_path, char *content_type)
 
 static int try_file(char *req_path, FILE **fp, char *content_type)
 {
+    // req_path error check 0
+    if(strlen(req_path) > REQ_PATH_CAP) {
+        err("try_file: requested path longer than %d characters", REQ_PATH_CAP);
+        return 1;
+    }
+
+    // req_path error check 1
     for(size_t i = 1; i < strlen(req_path); i++) {
         if(req_path[i-1] == '.' && req_path[i] == '.') {
             err("try_file: the requested path %s includes ..", req_path);
@@ -90,7 +103,7 @@ static int try_file(char *req_path, FILE **fp, char *content_type)
     }
 
     char file_path[COMMON_CAP];
-    sprintf(file_path, "%s%s", FILES,
+    snprintf(file_path, COMMON_CAP, "%s%s", FILES,
             (strlen(req_path) == 1) ? "/index.html" : req_path);
 
     *fp = fopen(file_path, "r");
@@ -190,14 +203,14 @@ static int handle_connection(sock_t *conn)
         return 1;
     }
 
-    // TODO: handle when the whole message is not sent
     ssize_t sent_sz = server_send(conn, send_buf, send_buf_sz);
     if(sent_sz < 0) {
         err("server_send: %s", strerror(errno));
         return 1;
+    } else if(sent_sz != send_buf_sz) {
+        err("server_send: sent %ld out of %ld bytes", sent_sz, send_buf_sz);
     }
 
-    info("server_send: sent %ld out of %ld bytes", sent_sz, send_buf_sz);
     return 0;
 }