Added userspace-based network timeout

The network in my whole city is down right now, so I've done some
error catching using timeouts etc without the kernel blocking
everything. Not having internet is exhausting though :(

2 files changed, 41 insertions, 9 deletions

diff --git a/libnet/dns.c b/libnet/dns.c
index fe57cc4..0ca2f99 100644
--- a/libnet/dns.c
+++ b/libnet/dns.c
@@ -96,7 +96,7 @@ static u32 dns_handle_packet(struct dns_packet *packet)
 u32 dns_request(const char *name)
 {
 	struct socket *socket = net_open(S_UDP);
-	if (!socket || !net_connect(socket, dns_ip_addr, 53) || part_count(name) == 1)
+	if (!socket || !net_connect(socket, dns_ip_addr, 53, NET_TIMEOUT) || part_count(name) == 1)
 		return 0;
 
 	u32 length = sizeof(struct dns_packet) + strlen(name) + part_count(name) + 4;
@@ -107,7 +107,7 @@ u32 dns_request(const char *name)
 	free(packet);
 
 	u8 buf[1024] = { 0 };
-	int l = net_receive(socket, buf, 1024);
+	int l = net_receive(socket, buf, 1024, NET_TIMEOUT);
 	net_close(socket);
 	if (l > 0)
 		return dns_handle_packet((void *)buf);
diff --git a/libnet/inc/net.h b/libnet/inc/net.h
index 5dd4ee1..6884690 100644
--- a/libnet/inc/net.h
+++ b/libnet/inc/net.h
@@ -6,6 +6,7 @@
 #include <dns.h>
 #include <http.h>
 #include <ip.h>
+#include <print.h>
 #include <socket.h>
 #include <sys.h>
 
@@ -18,10 +19,32 @@
 #define ip(a, b, c, d)                                                                             \
 	((((a)&0xff) << 24) | (((b)&0xff) << 16) | (((c)&0xff) << 8) | (((d)&0xff) << 0))
 
+#define NET_TIMEOUT 500
+#define NET_NO_TIMEOUT 0
+
+static inline int net_data_available(struct socket *socket)
+{
+	return (socket && socket->packets && socket->packets->head && socket->packets->head->data &&
+		((struct socket_data *)socket->packets->head->data)->length > 0);
+}
+
 #define net_open(type) (void *)sys1(SYS_NET_OPEN, (int)(type))
-#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))
+
+static inline int net_connect(struct socket *socket, u32 ip_addr, u16 dst_port, u32 timeout)
+{
+	if (!socket || !ip_addr || !dst_port)
+		return 0;
+	sys3(SYS_NET_CONNECT, (int)(socket), (int)(ip_addr), (int)(dst_port));
+	int time = time();
+	while (socket->state != S_CONNECTED) {
+		if (socket->state == S_FAILED || (timeout && time() - time >= timeout))
+			return 0;
+		yield();
+	}
+	return 1;
+}
+
 static inline int net_close(struct socket *socket)
 {
 	int res = 0;
@@ -29,12 +52,21 @@ static inline int net_close(struct socket *socket)
 		;
 	return res;
 }
-static inline int net_receive(struct socket *socket, void *buf, u32 len)
+
+static inline int net_receive(struct socket *socket, void *buf, u32 len, u32 timeout)
 {
-	int res = 0;
-	while (!(res = (int)sys3(SYS_NET_RECEIVE, (int)(socket), (int)(buf), (int)(len))))
-		;
-	return res;
+	if (!socket || !buf || !len)
+		return 0;
+
+	int time = time();
+	while (!net_data_available(socket)) {
+		if (socket->state == S_FAILED || (timeout && time() - time >= timeout))
+			return 0;
+		yield();
+	}
+
+	// TODO: Only return once all segments are received?
+	return (int)sys3(SYS_NET_RECEIVE, (int)(socket), (int)(buf), (int)(len));
 }
 
 #endif