Fixed some network race conditions

7 files changed, 46 insertions, 21 deletions

diff --git a/apps/browser.c b/apps/browser.c
index 2bc85d6..d03d963 100644
--- a/apps/browser.c
+++ b/apps/browser.c
@@ -89,7 +89,7 @@ void on_submit(void *event, struct element *box)
 	}
 	gui_sync(root, output);
 	gui_sync(root, code_label);
-	/* net_close(socket); */ // TODO: Fix net close before FIN/ACK got ACK'ed
+	net_close(socket);
 }
 
 int main()
diff --git a/kernel/features/net.c b/kernel/features/net.c
index 257d52b..436ad88 100644
--- a/kernel/features/net.c
+++ b/kernel/features/net.c
@@ -82,13 +82,13 @@ static int socket_close(struct socket *socket)
 	while (iterator != NULL && iterator->data != NULL) {
 		iterator = iterator->next;
 		if (iterator->data == socket)
-			return list_remove(list, iterator) != NULL;
+			list_remove(list, iterator);
 	}
 
 	socket->state = S_CLOSED;
 	free(socket);
 
-	return 0;
+	return 1;
 }
 
 /**
@@ -262,8 +262,9 @@ static void ip_send_packet(u32 dst, void *data, int len, u8 prot)
 	int arp_sent = 3;
 	u8 zero_hardware_addr[] = { 0, 0, 0, 0, 0, 0 };
 	u8 dst_mac[6];
-	sti();
-	if (same_net(dst))
+	if (same_net(dst)) {
+		scheduler_disable();
+		sti();
 		while (!arp_lookup(dst_mac, dst)) {
 			if (arp_sent) {
 				arp_sent--;
@@ -272,10 +273,12 @@ static void ip_send_packet(u32 dst, void *data, int len, u8 prot)
 				break;
 			}
 		}
-	else
+		cli();
+		scheduler_enable();
+	} else {
 		memcpy(dst_mac, gateway_mac, 6);
+	}
 
-	cli();
 	printf("Destination: %x:%x:%x:%x:%x:%x\n", dst_mac[0], dst_mac[1], dst_mac[2], dst_mac[3],
 	       dst_mac[4], dst_mac[5]);
 	ethernet_send_packet(dst_mac, (u8 *)packet, htons(packet->length), ETHERNET_TYPE_IP4);
@@ -488,6 +491,7 @@ static void tcp_handle_packet(struct tcp_packet *packet, u32 dst, int len)
 		tcp->ack_no = recv_seq;
 		tcp->seq_no = recv_ack;
 
+		socket->state = S_CONNECTED;
 		tcp->state++;
 		return;
 	} else if (tcp->state == 6 && (flags & 0xff) == (TCP_FLAG_ACK | TCP_FLAG_PSH)) {
@@ -508,6 +512,7 @@ static void tcp_handle_packet(struct tcp_packet *packet, u32 dst, int len)
 		tcp_send_packet(socket, TCP_FLAG_ACK, NULL, 0);
 		tcp_send_packet(socket, TCP_FLAG_FIN | TCP_FLAG_ACK, NULL, 0);
 
+		socket->state = S_CONNECTED;
 		tcp->state++;
 		return;
 	} else if (tcp->state == 7 && (flags & 0xff) == (TCP_FLAG_ACK | TCP_FLAG_FIN)) {
@@ -516,6 +521,8 @@ static void tcp_handle_packet(struct tcp_packet *packet, u32 dst, int len)
 
 		tcp_send_packet(socket, TCP_FLAG_ACK, NULL, 0);
 
+		proc_from_pid(socket->pid)->state = PROC_RUNNING;
+
 		socket->state = S_CLOSED;
 		tcp->state = 0;
 		return;
@@ -540,7 +547,6 @@ static void udp_handle_packet(struct udp_packet *packet)
 		return;
 	}
 
-	// TODO: Socket event to process
 	struct socket_data *sdata = malloc(sizeof(*sdata));
 	sdata->length = ntohs(packet->length);
 	if (sdata->length) {
@@ -551,6 +557,7 @@ static void udp_handle_packet(struct udp_packet *packet)
 	}
 	list_add(socket->packets, sdata);
 	proc_from_pid(socket->pid)->state = PROC_RUNNING;
+	/* printf("Waking up %d\n", socket->pid); */
 }
 
 static void ip_handle_packet(struct ip_packet *packet, int len)
@@ -714,10 +721,9 @@ struct socket *net_open(enum socket_type type)
 	return socket;
 }
 
-void net_close(struct socket *socket)
+int net_close(struct socket *socket)
 {
-	if (socket)
-		socket_close(socket);
+	return socket_close(socket);
 }
 
 int net_connect(struct socket *socket, u32 ip_addr, u16 dst_port)
@@ -738,11 +744,13 @@ int net_connect(struct socket *socket, u32 ip_addr, u16 dst_port)
 		socket->state = S_CONNECTING;
 		// TODO: Don't block kernel
 		tcp_send_packet(socket, TCP_FLAG_SYN, NULL, 0);
+		scheduler_disable();
 		sti();
 		u32 time = timer_get();
 		while (socket->state != S_CONNECTED && timer_get() - time < 1000)
 			;
 		cli();
+		scheduler_enable();
 		if (socket->state != S_CONNECTED) {
 			socket->state = S_FAILED;
 			return 0;
@@ -791,6 +799,7 @@ int net_receive(struct socket *socket, void *buf, u32 len)
 	while (iterator != NULL) {
 		struct socket_data *sdata = iterator->data;
 		u32 length = sdata->length;
+		// TODO: Fix underflow breaking
 		if (offset + length < len) {
 			memcpy((u8 *)buf + offset, sdata->data, length);
 			offset += length;
diff --git a/kernel/features/proc.c b/kernel/features/proc.c
index ffe0af0..f4e6124 100644
--- a/kernel/features/proc.c
+++ b/kernel/features/proc.c
@@ -67,12 +67,14 @@ void scheduler(struct regs *regs)
 
 void scheduler_enable(void)
 {
+	timer_install();
 	irq_install_handler(0, scheduler);
 }
 
 void scheduler_disable(void)
 {
-	irq_install_handler(0, scheduler);
+	irq_uninstall_handler(0);
+	timer_install();
 }
 
 void proc_print(void)
@@ -82,7 +84,8 @@ void proc_print(void)
 	printf("\nPROCESSES\n");
 	struct proc *proc = NULL;
 	while (node && (proc = node->data)) {
-		printf("Process %d: %s\n", proc->pid, proc->name);
+		printf("Process %d: %s [%s]\n", proc->pid, proc->name,
+		       proc->state == PROC_RUNNING ? "RUNNING" : "SLEEPING");
 		node = node->next;
 	}
 	printf("\n");
diff --git a/kernel/features/syscall.c b/kernel/features/syscall.c
index 3ee5a82..0328f47 100644
--- a/kernel/features/syscall.c
+++ b/kernel/features/syscall.c
@@ -97,7 +97,13 @@ void syscall_handler(struct regs *r)
 		break;
 	}
 	case SYS_NET_CLOSE: {
-		net_close((void *)r->ebx);
+		struct socket *s = (void *)r->ebx;
+		if (s->type == S_TCP && s->state != S_CLOSED) {
+			proc_current()->state = PROC_SLEEPING;
+			proc_yield(r);
+			return;
+		}
+		r->eax = net_close(s);
 		break;
 	}
 	case SYS_NET_CONNECT: {
@@ -112,6 +118,7 @@ void syscall_handler(struct regs *r)
 		if (!net_data_available((void *)r->ebx)) {
 			proc_current()->state = PROC_SLEEPING;
 			proc_yield(r);
+			return;
 		}
 		r->eax = net_receive((void *)r->ebx, (void *)r->ecx, r->edx);
 		break;
diff --git a/kernel/inc/net.h b/kernel/inc/net.h
index 0c36691..ddc2405 100644
--- a/kernel/inc/net.h
+++ b/kernel/inc/net.h
@@ -157,7 +157,7 @@ struct arp_table_entry {
 void ethernet_handle_packet(struct ethernet_packet *packet, int len);
 
 struct socket *net_open(enum socket_type type);
-void net_close(struct socket *socket);
+int net_close(struct socket *socket);
 int net_connect(struct socket *socket, u32 ip_addr, u16 dst_port);
 void net_send(struct socket *socket, void *data, u32 len);
 int net_receive(struct socket *socket, void *buf, u32 len);
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;
 }