Add peek option to NetSocket recv_from.

paul-marechal · Feb 16, 2020 · 2811f07 · 2811f07
1 parent cb68d83
commit 2811f07
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/core/io/net_socket.h b/core/io/net_socket.h
@@ -61,7 +61,7 @@ class NetSocket : public Reference {
 	virtual Error connect_to_host(IP_Address p_addr, uint16_t p_port) = 0;
 	virtual Error poll(PollType p_type, int timeout) const = 0;
 	virtual Error recv(uint8_t *p_buffer, int p_len, int &r_read) = 0;
-	virtual Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port) = 0;
+	virtual Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port, bool p_peek = false) = 0;
 	virtual Error send(const uint8_t *p_buffer, int p_len, int &r_sent) = 0;
 	virtual Error sendto(const uint8_t *p_buffer, int p_len, int &r_sent, IP_Address p_ip, uint16_t p_port) = 0;
 	virtual Ref<NetSocket> accept(IP_Address &r_ip, uint16_t &r_port) = 0;

diff --git a/drivers/unix/net_socket_posix.cpp b/drivers/unix/net_socket_posix.cpp
@@ -544,14 +544,14 @@ Error NetSocketPosix::recv(uint8_t *p_buffer, int p_len, int &r_read) {
 	return OK;
 }
 
-Error NetSocketPosix::recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port) {
+Error NetSocketPosix::recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port, bool p_peek) {
 	ERR_FAIL_COND_V(!is_open(), ERR_UNCONFIGURED);
 
 	struct sockaddr_storage from;
 	socklen_t len = sizeof(struct sockaddr_storage);
 	memset(&from, 0, len);
 
-	r_read = ::recvfrom(_sock, SOCK_BUF(p_buffer), p_len, 0, (struct sockaddr *)&from, &len);
+	r_read = ::recvfrom(_sock, SOCK_BUF(p_buffer), p_len, p_peek ? MSG_PEEK : 0, (struct sockaddr *)&from, &len);
 
 	if (r_read < 0) {
 		NetError err = _get_socket_error();

diff --git a/drivers/unix/net_socket_posix.h b/drivers/unix/net_socket_posix.h
@@ -81,7 +81,7 @@ class NetSocketPosix : public NetSocket {
 	virtual Error connect_to_host(IP_Address p_host, uint16_t p_port);
 	virtual Error poll(PollType p_type, int timeout) const;
 	virtual Error recv(uint8_t *p_buffer, int p_len, int &r_read);
-	virtual Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port);
+	virtual Error recvfrom(uint8_t *p_buffer, int p_len, int &r_read, IP_Address &r_ip, uint16_t &r_port, bool p_peek = false);
 	virtual Error send(const uint8_t *p_buffer, int p_len, int &r_sent);
 	virtual Error sendto(const uint8_t *p_buffer, int p_len, int &r_sent, IP_Address p_ip, uint16_t p_port);
 	virtual Ref<NetSocket> accept(IP_Address &r_ip, uint16_t &r_port);