From b6d229d726606ffe4af6e989b8ba55773af7ff19 Mon Sep 17 00:00:00 2001
From: Andy Green <andy@warmcat.com>
Date: Thu, 12 May 2016 21:04:33 +0800
Subject: [PATCH] check oom on lws_malloc

Signed-off-by: Andy Green <andy@warmcat.com>
---
 lib/libuv.c   |  4 ++++
 lib/rewrite.c |  5 +++++
 lib/server.c  | 53 ++++++++++++++++++++++++++++++---------------------
 lib/service.c |  2 ++
 4 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/lib/libuv.c b/lib/libuv.c
index a6234c70..7216078e 100644
--- a/lib/libuv.c
+++ b/lib/libuv.c
@@ -152,6 +152,10 @@ lws_uv_initloop(struct lws_context *context, uv_loop_t *loop, int tsi)
 
 	if (!loop) {
 		loop = lws_malloc(sizeof(*loop));
+		if (!loop) {
+			lwsl_err("OOM\n");
+			return -1;
+		}
 #if UV_VERSION_MAJOR > 0
 		uv_loop_init(loop);
 #else
diff --git a/lib/rewrite.c b/lib/rewrite.c
index db98e556..60a813d6 100644
--- a/lib/rewrite.c
+++ b/lib/rewrite.c
@@ -6,6 +6,11 @@ lws_rewrite_create(struct lws *wsi, hubbub_callback_t cb, const char *from, cons
 {
 	struct lws_rewrite *r = lws_malloc(sizeof(*r));
 
+	if (!r) {
+		lwsl_err("OOM\n");
+		return NULL;
+	}
+
 	if (hubbub_parser_create("UTF-8", false, &r->parser) != HUBBUB_OK) {
 		lws_free(r);
 
diff --git a/lib/server.c b/lib/server.c
index 87b580f5..6a921a3d 100644
--- a/lib/server.c
+++ b/lib/server.c
@@ -538,34 +538,39 @@ lws_http_action(struct lws *wsi)
 			lws_access_log(wsi);
 
 		wsi->access_log.header_log = lws_malloc(l);
+		if (wsi->access_log.header_log) {
 
-		tmp = localtime(&t);
-		if (tmp)
-			strftime(da, sizeof(da), "%d/%b/%Y:%H:%M:%S %z", tmp);
-		else
-			strcpy(da, "01/Jan/1970:00:00:00 +0000");
+			tmp = localtime(&t);
+			if (tmp)
+				strftime(da, sizeof(da), "%d/%b/%Y:%H:%M:%S %z", tmp);
+			else
+				strcpy(da, "01/Jan/1970:00:00:00 +0000");
 
-		pa = lws_get_peer_simple(wsi, ads, sizeof(ads));
-		if (!pa)
-			pa = "(unknown)";
+			pa = lws_get_peer_simple(wsi, ads, sizeof(ads));
+			if (!pa)
+				pa = "(unknown)";
 
-		if (meth >= 0)
-			me = method_names[meth];
-		else
-			me = "unknown";
+			if (meth >= 0)
+				me = method_names[meth];
+			else
+				me = "unknown";
 
-		snprintf(wsi->access_log.header_log, l,
-			 "%s - - [%s] \"%s %s %s\"",
-			 pa, da, me, uri_ptr,
-			 hver[wsi->u.http.request_version]);
+			snprintf(wsi->access_log.header_log, l,
+				 "%s - - [%s] \"%s %s %s\"",
+				 pa, da, me, uri_ptr,
+				 hver[wsi->u.http.request_version]);
 
-		l = lws_hdr_total_length(wsi, WSI_TOKEN_HTTP_USER_AGENT);
-		if (l) {
-			wsi->access_log.user_agent = lws_malloc(l + 2);
-			lws_hdr_copy(wsi, wsi->access_log.user_agent,
-				     l + 1, WSI_TOKEN_HTTP_USER_AGENT);
+			l = lws_hdr_total_length(wsi, WSI_TOKEN_HTTP_USER_AGENT);
+			if (l) {
+				wsi->access_log.user_agent = lws_malloc(l + 2);
+				if (wsi->access_log.user_agent)
+					lws_hdr_copy(wsi, wsi->access_log.user_agent,
+							l + 1, WSI_TOKEN_HTTP_USER_AGENT);
+				else
+					lwsl_err("OOM getting user agent\n");
+			}
+			wsi->access_log_pending = 1;
 		}
-		wsi->access_log_pending = 1;
 	}
 #endif
 
@@ -1449,6 +1454,10 @@ lws_adopt_socket_readbuf(struct lws_context *context, lws_sockfd_type accept_fd,
 	 * below to the rx buffer (via lws_header_table_reset()).
 	 */
 	wsi->u.hdr.preamble_rx = lws_malloc(len);
+	if (!wsi->u.hdr.preamble_rx) {
+		lwsl_err("OOM\n");
+		goto bail;
+	}
 	memcpy(wsi->u.hdr.preamble_rx, readbuf, len);
 	wsi->u.hdr.preamble_rx_len = len;
 
diff --git a/lib/service.c b/lib/service.c
index fcafd99d..03d2f5aa 100644
--- a/lib/service.c
+++ b/lib/service.c
@@ -368,6 +368,8 @@ int lws_rxflow_cache(struct lws *wsi, unsigned char *buf, int n, int len)
 	/* a new rxflow, buffer it and warn caller */
 	lwsl_info("new rxflow input buffer len %d\n", len - n);
 	wsi->rxflow_buffer = lws_malloc(len - n);
+	if (!wsi->rxflow_buffer)
+		return -1;
 	wsi->rxflow_len = len - n;
 	wsi->rxflow_pos = 0;
 	memcpy(wsi->rxflow_buffer, buf + n, len - n);