Skip to content

Commit c72b1a8

Browse files
scaprilescaprile
committed
add OOM handling to apps
1 parent 81b8799 commit c72b1a8

File tree

6 files changed

+240
-152
lines changed

6 files changed

+240
-152
lines changed

‎mongoose.c‎

Lines changed: 119 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1793,14 +1793,14 @@ int mg_http_parse(const char *s, size_t len, struct mg_http_message *hm) {
17931793
static void mg_http_vprintf_chunk(struct mg_connection *c, const char *fmt,
17941794
va_list *ap) {
17951795
size_t len = c->send.len;
1796-
mg_send(c, " \r\n", 10);
1796+
if (!mg_send(c, " \r\n", 10)) mg_error(c, "OOM");
17971797
mg_vxprintf(mg_pfn_iobuf, &c->send, fmt, ap);
17981798
if (c->send.len >= len + 10) {
17991799
mg_snprintf((char *) c->send.buf + len, 9, "%08lx", c->send.len - len - 10);
18001800
c->send.buf[len + 8] = '\r';
18011801
if (c->send.len == len + 10) c->is_resp = 0; // Last chunk, reset marker
18021802
}
1803-
mg_send(c, "\r\n", 2);
1803+
if (!mg_send(c, "\r\n", 2)) mg_error(c, "OOM");
18041804
}
18051805

18061806
void mg_http_printf_chunk(struct mg_connection *c, const char *fmt, ...) {
@@ -1812,8 +1812,7 @@ void mg_http_printf_chunk(struct mg_connection *c, const char *fmt, ...) {
18121812

18131813
void mg_http_write_chunk(struct mg_connection *c, const char *buf, size_t len) {
18141814
mg_printf(c, "%lx\r\n", (unsigned long) len);
1815-
mg_send(c, buf, len);
1816-
mg_send(c, "\r\n", 2);
1815+
if (!mg_send(c, buf, len) || !mg_send(c, "\r\n", 2)) mg_error(c, "OOM");
18171816
if (len == 0) c->is_resp = 0;
18181817
}
18191818

@@ -3353,8 +3352,8 @@ static const struct mg_mqtt_pmap s_prop_map[] = {
33533352
{MQTT_PROP_SUBSCRIPTION_IDENTIFIER_AVAILABLE, MQTT_PROP_TYPE_BYTE},
33543353
{MQTT_PROP_SHARED_SUBSCRIPTION_AVAILABLE, MQTT_PROP_TYPE_BYTE}};
33553354

3356-
void mg_mqtt_send_header(struct mg_connection *c, uint8_t cmd, uint8_t flags,
3357-
uint32_t len) {
3355+
static bool mqtt_send_header(struct mg_connection *c, uint8_t cmd,
3356+
uint8_t flags, uint32_t len) {
33583357
uint8_t buf[1 + sizeof(len)], *vlen = &buf[1];
33593358
buf[0] = (uint8_t) ((cmd << 4) | flags);
33603359
do {
@@ -3363,15 +3362,20 @@ void mg_mqtt_send_header(struct mg_connection *c, uint8_t cmd, uint8_t flags,
33633362
if (len > 0) *vlen |= 0x80;
33643363
vlen++;
33653364
} while (len > 0 && vlen < &buf[sizeof(buf)]);
3366-
mg_send(c, buf, (size_t) (vlen - buf));
3365+
return mg_send(c, buf, (size_t) (vlen - buf));
3366+
}
3367+
3368+
void mg_mqtt_send_header(struct mg_connection *c, uint8_t cmd, uint8_t flags,
3369+
uint32_t len) {
3370+
if (!mqtt_send_header(c, cmd, flags, len)) mg_error(c, "OOM");
33673371
}
33683372

3369-
static void mg_send_u16(struct mg_connection *c, uint16_t value) {
3370-
mg_send(c, &value, sizeof(value));
3373+
static bool mg_send_u16(struct mg_connection *c, uint16_t value) {
3374+
return mg_send(c, &value, sizeof(value));
33713375
}
33723376

3373-
static void mg_send_u32(struct mg_connection *c, uint32_t value) {
3374-
mg_send(c, &value, sizeof(value));
3377+
static bool mg_send_u32(struct mg_connection *c, uint32_t value) {
3378+
return mg_send(c, &value, sizeof(value));
33753379
}
33763380

33773381
static uint8_t varint_size(size_t length) {
@@ -3464,46 +3468,50 @@ static size_t get_props_size(struct mg_mqtt_prop *props, size_t count) {
34643468
return size;
34653469
}
34663470

3467-
static void mg_send_mqtt_properties(struct mg_connection *c,
3471+
static bool mg_send_mqtt_properties(struct mg_connection *c,
34683472
struct mg_mqtt_prop *props, size_t nprops) {
34693473
size_t total_size = get_properties_length(props, nprops);
34703474
uint8_t buf_v[4] = {0, 0, 0, 0};
34713475
uint8_t buf[4] = {0, 0, 0, 0};
34723476
size_t i, len = encode_varint(buf, total_size);
34733477

3474-
mg_send(c, buf, (size_t) len);
3478+
if (!mg_send(c, buf, (size_t) len)) return false;
34753479
for (i = 0; i < nprops; i++) {
3476-
mg_send(c, &props[i].id, sizeof(props[i].id));
3480+
if (!mg_send(c, &props[i].id, sizeof(props[i].id))) return false;
34773481
switch (mqtt_prop_type_by_id(props[i].id)) {
34783482
case MQTT_PROP_TYPE_STRING_PAIR:
3479-
mg_send_u16(c, mg_htons((uint16_t) props[i].key.len));
3480-
mg_send(c, props[i].key.buf, props[i].key.len);
3481-
mg_send_u16(c, mg_htons((uint16_t) props[i].val.len));
3482-
mg_send(c, props[i].val.buf, props[i].val.len);
3483+
if (!mg_send_u16(c, mg_htons((uint16_t) props[i].key.len)) ||
3484+
!mg_send(c, props[i].key.buf, props[i].key.len) ||
3485+
!mg_send_u16(c, mg_htons((uint16_t) props[i].val.len)) ||
3486+
!mg_send(c, props[i].val.buf, props[i].val.len))
3487+
return false;
34833488
break;
34843489
case MQTT_PROP_TYPE_BYTE:
3485-
mg_send(c, &props[i].iv, sizeof(uint8_t));
3490+
if (!mg_send(c, &props[i].iv, sizeof(uint8_t))) return false;
34863491
break;
34873492
case MQTT_PROP_TYPE_SHORT:
3488-
mg_send_u16(c, mg_htons((uint16_t) props[i].iv));
3493+
if (!mg_send_u16(c, mg_htons((uint16_t) props[i].iv))) return false;
34893494
break;
34903495
case MQTT_PROP_TYPE_INT:
3491-
mg_send_u32(c, mg_htonl((uint32_t) props[i].iv));
3496+
if (!mg_send_u32(c, mg_htonl((uint32_t) props[i].iv))) return false;
34923497
break;
34933498
case MQTT_PROP_TYPE_STRING:
3494-
mg_send_u16(c, mg_htons((uint16_t) props[i].val.len));
3495-
mg_send(c, props[i].val.buf, props[i].val.len);
3499+
if (!mg_send_u16(c, mg_htons((uint16_t) props[i].val.len)) ||
3500+
!mg_send(c, props[i].val.buf, props[i].val.len))
3501+
return false;
34963502
break;
34973503
case MQTT_PROP_TYPE_BINARY_DATA:
3498-
mg_send_u16(c, mg_htons((uint16_t) props[i].val.len));
3499-
mg_send(c, props[i].val.buf, props[i].val.len);
3504+
if (!mg_send_u16(c, mg_htons((uint16_t) props[i].val.len)) ||
3505+
!mg_send(c, props[i].val.buf, props[i].val.len))
3506+
return false;
35003507
break;
35013508
case MQTT_PROP_TYPE_VARIABLE_INT:
35023509
len = encode_varint(buf_v, props[i].iv);
3503-
mg_send(c, buf_v, (size_t) len);
3510+
if (!mg_send(c, buf_v, (size_t) len)) return false;
35043511
break;
35053512
}
35063513
}
3514+
return true;
35073515
}
35083516

35093517
size_t mg_mqtt_next_prop(struct mg_mqtt_message *msg, struct mg_mqtt_prop *prop,
@@ -3596,33 +3604,41 @@ void mg_mqtt_login(struct mg_connection *c, const struct mg_mqtt_opts *opts) {
35963604
total_len += get_props_size(opts->will_props, opts->num_will_props);
35973605
}
35983606

3599-
mg_mqtt_send_header(c, MQTT_CMD_CONNECT, 0, (uint32_t) total_len);
3600-
mg_send(c, hdr, sizeof(hdr));
36013607
// keepalive == 0 means "do not disconnect us!"
3602-
mg_send_u16(c, mg_htons((uint16_t) opts->keepalive));
3608+
if (!mqtt_send_header(c, MQTT_CMD_CONNECT, 0, (uint32_t) total_len) ||
3609+
!mg_send(c, hdr, sizeof(hdr)) ||
3610+
!mg_send_u16(c, mg_htons((uint16_t) opts->keepalive)))
3611+
goto fail;
36033612

3604-
if (c->is_mqtt5) mg_send_mqtt_properties(c, opts->props, opts->num_props);
3613+
if (c->is_mqtt5 && !mg_send_mqtt_properties(c, opts->props, opts->num_props))
3614+
goto fail;
36053615

3606-
mg_send_u16(c, mg_htons((uint16_t) cid.len));
3607-
mg_send(c, cid.buf, cid.len);
3616+
if (!mg_send_u16(c, mg_htons((uint16_t) cid.len)) ||
3617+
!mg_send(c, cid.buf, cid.len))
3618+
goto fail;
36083619

36093620
if (hdr[7] & MQTT_HAS_WILL) {
3610-
if (c->is_mqtt5)
3611-
mg_send_mqtt_properties(c, opts->will_props, opts->num_will_props);
3621+
if (c->is_mqtt5 &&
3622+
!mg_send_mqtt_properties(c, opts->will_props, opts->num_will_props))
3623+
goto fail;
36123624

3613-
mg_send_u16(c, mg_htons((uint16_t) opts->topic.len));
3614-
mg_send(c, opts->topic.buf, opts->topic.len);
3615-
mg_send_u16(c, mg_htons((uint16_t) opts->message.len));
3616-
mg_send(c, opts->message.buf, opts->message.len);
3617-
}
3618-
if (opts->user.len > 0) {
3619-
mg_send_u16(c, mg_htons((uint16_t) opts->user.len));
3620-
mg_send(c, opts->user.buf, opts->user.len);
3621-
}
3622-
if (opts->pass.len > 0) {
3623-
mg_send_u16(c, mg_htons((uint16_t) opts->pass.len));
3624-
mg_send(c, opts->pass.buf, opts->pass.len);
3625+
if (!mg_send_u16(c, mg_htons((uint16_t) opts->topic.len)) ||
3626+
!mg_send(c, opts->topic.buf, opts->topic.len) ||
3627+
!mg_send_u16(c, mg_htons((uint16_t) opts->message.len)) ||
3628+
!mg_send(c, opts->message.buf, opts->message.len))
3629+
goto fail;
36253630
}
3631+
if (opts->user.len > 0 &&
3632+
(!mg_send_u16(c, mg_htons((uint16_t) opts->user.len)) ||
3633+
!mg_send(c, opts->user.buf, opts->user.len)))
3634+
goto fail;
3635+
if (opts->pass.len > 0 &&
3636+
(!mg_send_u16(c, mg_htons((uint16_t) opts->pass.len)) ||
3637+
!mg_send(c, opts->pass.buf, opts->pass.len)))
3638+
goto fail;
3639+
return;
3640+
fail:
3641+
mg_error(c, "OOM");
36263642
}
36273643

36283644
uint16_t mg_mqtt_pub(struct mg_connection *c, const struct mg_mqtt_opts *opts) {
@@ -3637,20 +3653,28 @@ uint16_t mg_mqtt_pub(struct mg_connection *c, const struct mg_mqtt_opts *opts) {
36373653
if (c->is_mqtt5) len += get_props_size(opts->props, opts->num_props);
36383654

36393655
if (opts->qos > 0 && id != 0) flags |= 1 << 3;
3640-
mg_mqtt_send_header(c, MQTT_CMD_PUBLISH, flags, (uint32_t) len);
3641-
mg_send_u16(c, mg_htons((uint16_t) opts->topic.len));
3642-
mg_send(c, opts->topic.buf, opts->topic.len);
3656+
if (!mqtt_send_header(c, MQTT_CMD_PUBLISH, flags, (uint32_t) len) ||
3657+
!mg_send_u16(c, mg_htons((uint16_t) opts->topic.len)) ||
3658+
!mg_send(c, opts->topic.buf, opts->topic.len))
3659+
goto fail;
36433660
if (opts->qos > 0) { // need to send 'id' field
36443661
if (id == 0) { // generate new one if not resending
36453662
if (++c->mgr->mqtt_id == 0) ++c->mgr->mqtt_id;
36463663
id = c->mgr->mqtt_id;
36473664
}
3648-
mg_send_u16(c, mg_htons(id));
3665+
if (!mg_send_u16(c, mg_htons(id))) goto fail;
36493666
}
36503667

3651-
if (c->is_mqtt5) mg_send_mqtt_properties(c, opts->props, opts->num_props);
3668+
if (c->is_mqtt5 && !mg_send_mqtt_properties(c, opts->props, opts->num_props))
3669+
goto fail;
3670+
3671+
if (opts->message.len > 0 &&
3672+
!mg_send(c, opts->message.buf, opts->message.len))
3673+
goto fail;
3674+
return id;
36523675

3653-
if (opts->message.len > 0) mg_send(c, opts->message.buf, opts->message.len);
3676+
fail:
3677+
mg_error(c, "OOM");
36543678
return id;
36553679
}
36563680

@@ -3659,14 +3683,20 @@ void mg_mqtt_sub(struct mg_connection *c, const struct mg_mqtt_opts *opts) {
36593683
size_t plen = c->is_mqtt5 ? get_props_size(opts->props, opts->num_props) : 0;
36603684
size_t len = 2 + opts->topic.len + 2 + 1 + plen;
36613685

3662-
mg_mqtt_send_header(c, MQTT_CMD_SUBSCRIBE, 2, (uint32_t) len);
3686+
if (!mqtt_send_header(c, MQTT_CMD_SUBSCRIBE, 2, (uint32_t) len)) goto fail;
36633687
if (++c->mgr->mqtt_id == 0) ++c->mgr->mqtt_id;
3664-
mg_send_u16(c, mg_htons(c->mgr->mqtt_id));
3665-
if (c->is_mqtt5) mg_send_mqtt_properties(c, opts->props, opts->num_props);
3688+
if (!mg_send_u16(c, mg_htons(c->mgr->mqtt_id))) goto fail;
3689+
3690+
if (c->is_mqtt5 && !mg_send_mqtt_properties(c, opts->props, opts->num_props))
3691+
goto fail;
36663692

3667-
mg_send_u16(c, mg_htons((uint16_t) opts->topic.len));
3668-
mg_send(c, opts->topic.buf, opts->topic.len);
3669-
mg_send(c, &qos_, sizeof(qos_));
3693+
if (!mg_send_u16(c, mg_htons((uint16_t) opts->topic.len)) ||
3694+
!mg_send(c, opts->topic.buf, opts->topic.len) ||
3695+
!mg_send(c, &qos_, sizeof(qos_)))
3696+
goto fail;
3697+
return;
3698+
fail:
3699+
mg_error(c, "OOM");
36703700
}
36713701

36723702
int mg_mqtt_parse(const uint8_t *buf, size_t len, uint8_t version,
@@ -3773,15 +3803,16 @@ static void mqtt_cb(struct mg_connection *c, int ev, void *ev_data) {
37733803
uint32_t remaining_len = sizeof(id);
37743804
if (c->is_mqtt5) remaining_len += 2; // 3.4.2
37753805

3776-
mg_mqtt_send_header(
3777-
c,
3778-
(uint8_t) (mm.qos == 2 ? MQTT_CMD_PUBREC : MQTT_CMD_PUBACK),
3779-
0, remaining_len);
3780-
mg_send(c, &id, sizeof(id));
3806+
if (!mqtt_send_header(c,
3807+
(uint8_t) (mm.qos == 2 ? MQTT_CMD_PUBREC
3808+
: MQTT_CMD_PUBACK),
3809+
0, remaining_len) ||
3810+
!mg_send(c, &id, sizeof(id)))
3811+
goto fail;
37813812

37823813
if (c->is_mqtt5) {
37833814
uint16_t zero = 0;
3784-
mg_send(c, &zero, sizeof(zero));
3815+
if (!mg_send(c, &zero, sizeof(zero))) goto fail;
37853816
}
37863817
}
37873818
mg_call(c, MG_EV_MQTT_MSG, &mm); // let the app handle qos stuff
@@ -3790,15 +3821,18 @@ static void mqtt_cb(struct mg_connection *c, int ev, void *ev_data) {
37903821
case MQTT_CMD_PUBREC: { // MQTT5: 3.5.2-1 TODO(): variable header rc
37913822
uint16_t id = mg_ntohs(mm.id);
37923823
uint32_t remaining_len = sizeof(id); // MQTT5 3.6.2-1
3793-
mg_mqtt_send_header(c, MQTT_CMD_PUBREL, 2, remaining_len);
3794-
mg_send(c, &id, sizeof(id)); // MQTT5 3.6.1-1, flags = 2
3824+
if (!mqtt_send_header(c, MQTT_CMD_PUBREL, 2,
3825+
remaining_len) // MQTT5 3.6.1-1, flags = 2
3826+
|| !mg_send(c, &id, sizeof(id)))
3827+
goto fail;
37953828
break;
37963829
}
37973830
case MQTT_CMD_PUBREL: { // MQTT5: 3.6.2-1 TODO(): variable header rc
37983831
uint16_t id = mg_ntohs(mm.id);
37993832
uint32_t remaining_len = sizeof(id); // MQTT5 3.7.2-1
3800-
mg_mqtt_send_header(c, MQTT_CMD_PUBCOMP, 0, remaining_len);
3801-
mg_send(c, &id, sizeof(id));
3833+
if (!mqtt_send_header(c, MQTT_CMD_PUBCOMP, 0, remaining_len) ||
3834+
!mg_send(c, &id, sizeof(id)))
3835+
goto fail;
38023836
break;
38033837
}
38043838
}
@@ -3810,6 +3844,9 @@ static void mqtt_cb(struct mg_connection *c, int ev, void *ev_data) {
38103844
}
38113845
}
38123846
(void) ev_data;
3847+
return;
3848+
fail:
3849+
mg_error(c, "OOM");
38133850
}
38143851

38153852
void mg_mqtt_ping(struct mg_connection *nc) {
@@ -3824,19 +3861,24 @@ void mg_mqtt_disconnect(struct mg_connection *c,
38243861
const struct mg_mqtt_opts *opts) {
38253862
size_t len = 0;
38263863
if (c->is_mqtt5) len = 1 + get_props_size(opts->props, opts->num_props);
3827-
mg_mqtt_send_header(c, MQTT_CMD_DISCONNECT, 0, (uint32_t) len);
3864+
if (!mqtt_send_header(c, MQTT_CMD_DISCONNECT, 0, (uint32_t) len)) goto fail;
38283865

38293866
if (c->is_mqtt5) {
38303867
uint8_t zero = 0;
3831-
mg_send(c, &zero, sizeof(zero)); // reason code
3832-
mg_send_mqtt_properties(c, opts->props, opts->num_props);
3868+
if (!mg_send(c, &zero, sizeof(zero)) // reason code
3869+
|| !mg_send_mqtt_properties(c, opts->props, opts->num_props))
3870+
goto fail;
38333871
}
3872+
return;
3873+
fail:
3874+
mg_error(c, "OOM");
38343875
}
38353876

38363877
struct mg_connection *mg_mqtt_connect(struct mg_mgr *mgr, const char *url,
38373878
const struct mg_mqtt_opts *opts,
38383879
mg_event_handler_t fn, void *fn_data) {
3839-
struct mg_connection *c = mg_connect_svc(mgr, url, fn, fn_data, mqtt_cb, NULL);
3880+
struct mg_connection *c =
3881+
mg_connect_svc(mgr, url, fn, fn_data, mqtt_cb, NULL);
38403882
if (c != NULL) {
38413883
struct mg_mqtt_opts empty;
38423884
memset(&empty, 0, sizeof(empty));
@@ -5537,6 +5579,7 @@ bool mg_open_listener(struct mg_connection *c, const char *url) {
55375579
static void write_conn(struct mg_connection *c) {
55385580
long len = c->is_tls ? mg_tls_send(c, c->send.buf, c->send.len)
55395581
: mg_io_send(c, c->send.buf, c->send.len);
5582+
// TODO(): mg_tls_send() may return 0 forever on steady OOM
55405583
if (len == MG_IO_ERR) {
55415584
mg_error(c, "tx err");
55425585
} else if (len > 0) {
@@ -5615,7 +5658,7 @@ bool mg_send(struct mg_connection *c, const void *buf, size_t len) {
56155658
res = tx_udp(ifp, s->mac, ifp->ip, c->loc.port, rem_ip, c->rem.port, buf,
56165659
len);
56175660
} else {
5618-
res = mg_iobuf_add(&c->send, c->send.len, buf, len);
5661+
res = (bool) mg_iobuf_add(&c->send, c->send.len, buf, len);
56195662
// res == 0 means an OOM condition (iobuf couldn't resize), yet this is so
56205663
// far recoverable, let the caller decide
56215664
}
@@ -8972,7 +9015,7 @@ bool mg_send(struct mg_connection *c, const void *buf, size_t len) {
89729015
iolog(c, (char *) buf, n, false);
89739016
return n > 0;
89749017
} else {
8975-
return mg_iobuf_add(&c->send, c->send.len, buf, len);
9018+
return (bool) mg_iobuf_add(&c->send, c->send.len, buf, len);
89769019
// returning 0 means an OOM condition (iobuf couldn't resize), yet this is
89779020
// so far recoverable, let the caller decide
89789021
}
@@ -9172,6 +9215,7 @@ static void write_conn(struct mg_connection *c) {
91729215
char *buf = (char *) c->send.buf;
91739216
size_t len = c->send.len;
91749217
long n = c->is_tls ? mg_tls_send(c, buf, len) : mg_io_send(c, buf, len);
9218+
// TODO(): mg_tls_send() may return 0 forever on steady OOM
91759219
MG_DEBUG(("%lu %ld snd %ld/%ld rcv %ld/%ld n=%ld err=%d", c->id, c->fd,
91769220
(long) c->send.len, (long) c->send.size, (long) c->recv.len,
91779221
(long) c->recv.size, n, MG_SOCK_ERR(n)));
@@ -20349,7 +20393,7 @@ static void ws_handshake(struct mg_connection *c, const struct mg_str *wskey,
2034920393
mg_printf(c, "Sec-WebSocket-Protocol: %.*s\r\n", (int) wsproto->len,
2035020394
wsproto->buf);
2035120395
}
20352-
mg_send(c, "\r\n", 2);
20396+
if (!mg_send(c, "\r\n", 2)) mg_error(c, "OOM");
2035320397
}
2035420398

2035520399
static uint32_t be32(const uint8_t *p) {

0 commit comments

Comments
 (0)