Skip to content

Commit e844776

Browse files
projectgusdpgeorge
projectgus
authored and
dpgeorge
committed
extmod/modlwip: Add optional flags argument for recv and recvfrom.
Implements MSG_PEEK and MSG_DONTWAIT. This work was funded through GitHub Sponsors. Signed-off-by: Angus Gratton <angus@redyak.com.au>
1 parent 4dff9cb commit e844776

File tree

1 file changed

+62
-78
lines changed

1 file changed

+62
-78
lines changed

‎extmod/modlwip.c

Lines changed: 62 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@
7070

7171
#define TCP_NODELAY TF_NODELAY
7272

73+
// Socket flags
74+
#define MSG_PEEK 0x01
75+
#define MSG_DONTWAIT 0x02
76+
7377
// For compatibilily with older lwIP versions.
7478
#ifndef ip_set_option
7579
#define ip_set_option(pcb, opt) ((pcb)->so_options |= (opt))
@@ -673,13 +677,13 @@ static mp_uint_t lwip_raw_udp_send(lwip_socket_obj_t *socket, const byte *buf, m
673677
}
674678

675679
// Helper function for recv/recvfrom to handle raw/UDP packets
676-
static mp_uint_t lwip_raw_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, ip_addr_t *ip, mp_uint_t *port, int *_errno) {
680+
static mp_uint_t lwip_raw_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, mp_int_t flags, ip_addr_t *ip, mp_uint_t *port, int *_errno) {
677681

678682
lwip_incoming_packet_t *slot = &socket->incoming.udp_raw.array[socket->incoming.udp_raw.iget];
679683

680684
if (slot->pbuf == NULL) {
681-
if (socket->timeout == 0) {
682-
// Non-blocking socket.
685+
// Non-blocking socket or flag
686+
if (socket->timeout == 0 || (flags & MSG_DONTWAIT)) {
683687
*_errno = MP_EAGAIN;
684688
return -1;
685689
}
@@ -705,9 +709,11 @@ static mp_uint_t lwip_raw_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_u
705709
MICROPY_PY_LWIP_ENTER
706710

707711
u16_t result = pbuf_copy_partial(p, buf, ((p->tot_len > len) ? len : p->tot_len), 0);
708-
pbuf_free(p);
709-
slot->pbuf = NULL;
710-
socket->incoming.udp_raw.iget = (socket->incoming.udp_raw.iget + 1) % LWIP_INCOMING_PACKET_QUEUE_LEN;
712+
if ((flags & MSG_PEEK) == 0) {
713+
pbuf_free(p);
714+
slot->pbuf = NULL;
715+
socket->incoming.udp_raw.iget = (socket->incoming.udp_raw.iget + 1) % LWIP_INCOMING_PACKET_QUEUE_LEN;
716+
}
711717

712718
MICROPY_PY_LWIP_EXIT
713719

@@ -815,14 +821,14 @@ static mp_uint_t lwip_tcp_send(lwip_socket_obj_t *socket, const byte *buf, mp_ui
815821
}
816822

817823
// Helper function for recv/recvfrom to handle TCP packets
818-
static mp_uint_t lwip_tcp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, int *_errno) {
824+
static mp_uint_t lwip_tcp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, mp_int_t flags, int *_errno) {
819825
// Check for any pending errors
820826
STREAM_ERROR_CHECK(socket);
821827

822828
if (socket->incoming.tcp.pbuf == NULL) {
823829

824-
// Non-blocking socket
825-
if (socket->timeout == 0) {
830+
// Non-blocking socket or flag
831+
if (socket->timeout == 0 || (flags & MSG_DONTWAIT)) {
826832
if (socket->state == STATE_PEER_CLOSED) {
827833
return 0;
828834
}
@@ -867,19 +873,21 @@ static mp_uint_t lwip_tcp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_
867873

868874
memcpy(buf, (byte *)p->payload + socket->recv_offset, len);
869875

870-
remaining -= len;
871-
if (remaining == 0) {
872-
socket->incoming.tcp.pbuf = p->next;
873-
// If we don't ref here, free() will free the entire chain,
874-
// if we ref, it does what we need: frees 1st buf, and decrements
875-
// next buf's refcount back to 1.
876-
pbuf_ref(p->next);
877-
pbuf_free(p);
878-
socket->recv_offset = 0;
879-
} else {
880-
socket->recv_offset += len;
876+
if ((flags & MSG_PEEK) == 0) {
877+
remaining -= len;
878+
if (remaining == 0) {
879+
socket->incoming.tcp.pbuf = p->next;
880+
// If we don't ref here, free() will free the entire chain,
881+
// if we ref, it does what we need: frees 1st buf, and decrements
882+
// next buf's refcount back to 1.
883+
pbuf_ref(p->next);
884+
pbuf_free(p);
885+
socket->recv_offset = 0;
886+
} else {
887+
socket->recv_offset += len;
888+
}
889+
tcp_recved(socket->pcb.tcp, len);
881890
}
882-
tcp_recved(socket->pcb.tcp, len);
883891

884892
MICROPY_PY_LWIP_EXIT
885893

@@ -1271,40 +1279,58 @@ static mp_obj_t lwip_socket_send(mp_obj_t self_in, mp_obj_t buf_in) {
12711279
}
12721280
static MP_DEFINE_CONST_FUN_OBJ_2(lwip_socket_send_obj, lwip_socket_send);
12731281

1274-
static mp_obj_t lwip_socket_recv(mp_obj_t self_in, mp_obj_t len_in) {
1275-
lwip_socket_obj_t *socket = MP_OBJ_TO_PTR(self_in);
1282+
// Common implementation for recv & recvfrom
1283+
static mp_obj_t lwip_socket_recv_common(size_t n_args, const mp_obj_t *args, ip_addr_t *ip, mp_uint_t *port) {
1284+
lwip_socket_obj_t *socket = MP_OBJ_TO_PTR(args[0]);
1285+
mp_int_t len = mp_obj_get_int(args[1]);
1286+
mp_int_t flags = n_args > 2 ? mp_obj_get_int(args[2]) : 0;
12761287
int _errno;
1288+
vstr_t vstr;
1289+
mp_uint_t ret = 0;
12771290

12781291
lwip_socket_check_connected(socket);
12791292

1280-
mp_int_t len = mp_obj_get_int(len_in);
1281-
vstr_t vstr;
12821293
vstr_init_len(&vstr, len);
12831294

1284-
mp_uint_t ret = 0;
12851295
switch (socket->type) {
1286-
case MOD_NETWORK_SOCK_STREAM: {
1287-
ret = lwip_tcp_receive(socket, (byte *)vstr.buf, len, &_errno);
1296+
case MOD_NETWORK_SOCK_STREAM:
1297+
if (ip != NULL) {
1298+
*ip = socket->tcp_peer_addr;
1299+
*port = (mp_uint_t)socket->tcp_peer_port;
1300+
}
1301+
ret = lwip_tcp_receive(socket, (byte *)vstr.buf, len, flags, &_errno);
12881302
break;
1289-
}
12901303
case MOD_NETWORK_SOCK_DGRAM:
12911304
#if MICROPY_PY_LWIP_SOCK_RAW
12921305
case MOD_NETWORK_SOCK_RAW:
12931306
#endif
1294-
ret = lwip_raw_udp_receive(socket, (byte *)vstr.buf, len, NULL, NULL, &_errno);
1307+
ret = lwip_raw_udp_receive(socket, (byte *)vstr.buf, len, flags, ip, port, &_errno);
12951308
break;
12961309
}
12971310
if (ret == -1) {
12981311
mp_raise_OSError(_errno);
12991312
}
1300-
13011313
if (ret == 0) {
13021314
return mp_const_empty_bytes;
13031315
}
13041316
vstr.len = ret;
13051317
return mp_obj_new_bytes_from_vstr(&vstr);
13061318
}
1307-
static MP_DEFINE_CONST_FUN_OBJ_2(lwip_socket_recv_obj, lwip_socket_recv);
1319+
1320+
static mp_obj_t lwip_socket_recv(size_t n_args, const mp_obj_t *args) {
1321+
return lwip_socket_recv_common(n_args, args, NULL, NULL);
1322+
}
1323+
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(lwip_socket_recv_obj, 2, 3, lwip_socket_recv);
1324+
1325+
static mp_obj_t lwip_socket_recvfrom(size_t n_args, const mp_obj_t *args) {
1326+
ip_addr_t ip;
1327+
mp_uint_t port;
1328+
mp_obj_t tuple[2];
1329+
tuple[0] = lwip_socket_recv_common(n_args, args, &ip, &port);
1330+
tuple[1] = lwip_format_inet_addr(&ip, port);
1331+
return mp_obj_new_tuple(2, tuple);
1332+
}
1333+
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(lwip_socket_recvfrom_obj, 2, 3, lwip_socket_recvfrom);
13081334

13091335
static mp_obj_t lwip_socket_sendto(mp_obj_t self_in, mp_obj_t data_in, mp_obj_t addr_in) {
13101336
lwip_socket_obj_t *socket = MP_OBJ_TO_PTR(self_in);
@@ -1339,50 +1365,6 @@ static mp_obj_t lwip_socket_sendto(mp_obj_t self_in, mp_obj_t data_in, mp_obj_t
13391365
}
13401366
static MP_DEFINE_CONST_FUN_OBJ_3(lwip_socket_sendto_obj, lwip_socket_sendto);
13411367

1342-
static mp_obj_t lwip_socket_recvfrom(mp_obj_t self_in, mp_obj_t len_in) {
1343-
lwip_socket_obj_t *socket = MP_OBJ_TO_PTR(self_in);
1344-
int _errno;
1345-
1346-
lwip_socket_check_connected(socket);
1347-
1348-
mp_int_t len = mp_obj_get_int(len_in);
1349-
vstr_t vstr;
1350-
vstr_init_len(&vstr, len);
1351-
ip_addr_t ip;
1352-
mp_uint_t port;
1353-
1354-
mp_uint_t ret = 0;
1355-
switch (socket->type) {
1356-
case MOD_NETWORK_SOCK_STREAM: {
1357-
ip = socket->tcp_peer_addr;
1358-
port = (mp_uint_t)socket->tcp_peer_port;
1359-
ret = lwip_tcp_receive(socket, (byte *)vstr.buf, len, &_errno);
1360-
break;
1361-
}
1362-
case MOD_NETWORK_SOCK_DGRAM:
1363-
#if MICROPY_PY_LWIP_SOCK_RAW
1364-
case MOD_NETWORK_SOCK_RAW:
1365-
#endif
1366-
ret = lwip_raw_udp_receive(socket, (byte *)vstr.buf, len, &ip, &port, &_errno);
1367-
break;
1368-
}
1369-
if (ret == -1) {
1370-
mp_raise_OSError(_errno);
1371-
}
1372-
1373-
mp_obj_t tuple[2];
1374-
if (ret == 0) {
1375-
tuple[0] = mp_const_empty_bytes;
1376-
} else {
1377-
vstr.len = ret;
1378-
tuple[0] = mp_obj_new_bytes_from_vstr(&vstr);
1379-
}
1380-
tuple[1] = lwip_format_inet_addr(&ip, port);
1381-
1382-
return mp_obj_new_tuple(2, tuple);
1383-
}
1384-
static MP_DEFINE_CONST_FUN_OBJ_2(lwip_socket_recvfrom_obj, lwip_socket_recvfrom);
1385-
13861368
static mp_obj_t lwip_socket_sendall(mp_obj_t self_in, mp_obj_t buf_in) {
13871369
lwip_socket_obj_t *socket = MP_OBJ_TO_PTR(self_in);
13881370
lwip_socket_check_connected(socket);
@@ -1542,12 +1524,12 @@ static mp_uint_t lwip_socket_read(mp_obj_t self_in, void *buf, mp_uint_t size, i
15421524

15431525
switch (socket->type) {
15441526
case MOD_NETWORK_SOCK_STREAM:
1545-
return lwip_tcp_receive(socket, buf, size, errcode);
1527+
return lwip_tcp_receive(socket, buf, size, 0, errcode);
15461528
case MOD_NETWORK_SOCK_DGRAM:
15471529
#if MICROPY_PY_LWIP_SOCK_RAW
15481530
case MOD_NETWORK_SOCK_RAW:
15491531
#endif
1550-
return lwip_raw_udp_receive(socket, buf, size, NULL, NULL, errcode);
1532+
return lwip_raw_udp_receive(socket, buf, size, 0, NULL, NULL, errcode);
15511533
}
15521534
// Unreachable
15531535
return MP_STREAM_ERROR;
@@ -1919,6 +1901,8 @@ static const mp_rom_map_elem_t mp_module_lwip_globals_table[] = {
19191901

19201902
{ MP_ROM_QSTR(MP_QSTR_IPPROTO_TCP), MP_ROM_INT(IP_PROTO_TCP) },
19211903
{ MP_ROM_QSTR(MP_QSTR_TCP_NODELAY), MP_ROM_INT(TCP_NODELAY) },
1904+
{ MP_ROM_QSTR(MP_QSTR_MSG_PEEK), MP_ROM_INT(MSG_PEEK) },
1905+
{ MP_ROM_QSTR(MP_QSTR_MSG_DONTWAIT), MP_ROM_INT(MSG_DONTWAIT) },
19221906
};
19231907

19241908
static MP_DEFINE_CONST_DICT(mp_module_lwip_globals, mp_module_lwip_globals_table);

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy