Fixed most memory leaks in response.c

1 files changed, 16 insertions, 7 deletions

diff --git a/response.c b/response.c
index 959ed18..8866bd1 100644
--- a/response.c
+++ b/response.c
@@ -30,6 +30,7 @@ int response_send(lua_State * L, const int fd)
 	int bytes_total = 2048 * sizeof(char);
 	int bytes_used = 0;
 	char * buffer = malloc(bytes_total);
+	char * new_buffer = NULL;
 
 	lua_pushstring(L, "status");
 	lua_gettable(L, -2);
@@ -38,6 +39,7 @@ int response_send(lua_State * L, const int fd)
 
 	if (0 == status || NULL == buffer)
 	{
+		free(buffer);
 		lua_pop(L, 1);
 		return write(fd, error_response, strlen(error_response));
 	}
@@ -49,8 +51,10 @@ int response_send(lua_State * L, const int fd)
 
 	if (0 == lua_istable(L, -1))
 	{
+		const int result = write(fd, buffer, bytes_used + 2);
+		free(buffer);
 		lua_pop(L, 2);
-		return write(fd, buffer, bytes_used + 2);
+		return result;
 	}
 
 	static const char * header_pattern = "%s: %s\r\n";
@@ -71,12 +75,14 @@ int response_send(lua_State * L, const int fd)
 				lua_pop(L, 4);
 				return write(fd, error_response, strlen(error_response));
 			}
-			buffer = realloc(buffer, bytes_total);
-			if (NULL == buffer)
+			new_buffer = realloc(buffer, bytes_total); // TODO: Realloc and error handling involves duplication.
+			if (NULL == new_buffer)
 			{
+				free(buffer);
 				lua_pop(L, 4);
 				return write(fd, error_response, strlen(error_response));
 			}
+			buffer = new_buffer;
 			new_bytes = snprintf(buffer + bytes_used, bytes_left, header_pattern, key, value);
 		}
 		bytes_used += new_bytes;
@@ -86,12 +92,13 @@ int response_send(lua_State * L, const int fd)
 	if (bytes_total < bytes_used + 2)
 	{
 		bytes_total += 2;
-		buffer = realloc(buffer, bytes_total);
-		if (NULL == buffer)
+		new_buffer = realloc(buffer, bytes_total);
+		if (NULL == new_buffer)
 		{
 			lua_pop(L, 2);
 			return write(fd, error_response, strlen(error_response));
 		}
+		buffer = new_buffer;
 	}
 
 	lua_pop(L, 1);
@@ -110,12 +117,14 @@ int response_send(lua_State * L, const int fd)
 		if (bytes_total < bytes_used + (int) data_length)
 		{
 			bytes_total += data_length;
-			buffer = realloc(buffer, bytes_total);
-			if (NULL == buffer)
+			new_buffer = realloc(buffer, bytes_total);
+			if (NULL == new_buffer)
 			{
+				free(buffer);
 				lua_pop(L, 2);
 				return write(fd, error_response, strlen(error_response));
 			}
+			buffer = new_buffer;
 		}
 
 		memcpy(buffer + bytes_used, data, data_length);