Fixed some network race conditions

2 files changed, 12 insertions, 6 deletions

diff --git a/libnet/dns.c b/libnet/dns.c
index fc6b06f..6682831 100644
--- a/libnet/dns.c
+++ b/libnet/dns.c
@@ -59,8 +59,6 @@ static u32 dns_handle_packet(struct dns_packet *packet)
 u32 dns_request(const char *name, const char *tld)
 {
 	struct socket *socket = net_open(S_UDP);
-	/* if (socket) */
-	/* 	socket->src_port = 50053; */
 	if (!socket || !net_connect(socket, dns_ip_addr, 53))
 		return 0;
 
@@ -71,8 +69,9 @@ u32 dns_request(const char *name, const char *tld)
 	net_send(socket, packet, length);
 	free(packet);
 
-	u8 buf[128] = { 0 };
-	int l = net_receive(socket, buf, 128);
+	u8 buf[1024] = { 0 };
+	int l = net_receive(socket, buf, 1024);
+	net_close(socket);
 	if (l > 0)
 		return dns_handle_packet((void *)buf);
 	else
diff --git a/libnet/inc/net.h b/libnet/inc/net.h
index 6ff9619..ac9296e 100644
--- a/libnet/inc/net.h
+++ b/libnet/inc/net.h
@@ -28,14 +28,21 @@
 	((((a)&0xff) << 24) | (((b)&0xff) << 16) | (((c)&0xff) << 8) | (((d)&0xff) << 0))
 
 #define net_open(type) (void *)sys1(SYS_NET_OPEN, (int)(type))
-#define net_close(socket) (void)sys1(SYS_NET_CLOSE, (int)(socket))
 #define net_connect(socket, ip_addr, dst_port)                                                     \
 	(int)sys3(SYS_NET_CONNECT, (int)(socket), (int)(ip_addr), (int)(dst_port))
 #define net_send(socket, data, len) (void)sys3(SYS_NET_SEND, (int)(socket), (int)(data), (int)(len))
+#include <print.h>
+static inline int net_close(struct socket *socket)
+{
+	int res = 0;
+	while (!(res = (int)sys1(SYS_NET_CLOSE, (int)(socket))))
+		;
+	return res;
+}
 static inline int net_receive(struct socket *socket, void *buf, u32 len)
 {
 	int res = 0;
-	while ((res = (int)sys3(SYS_NET_RECEIVE, (int)(socket), (int)(buf), (int)(len))) == 0)
+	while (!(res = (int)sys3(SYS_NET_RECEIVE, (int)(socket), (int)(buf), (int)(len))))
 		;
 	return res;
 }