From 2ea9e035222ae677a9585ce60d406347c4a328ff Mon Sep 17 00:00:00 2001
From: ValKmjolnir <lhk101lhk101@qq.com>
Date: Mon, 6 Nov 2023 00:22:11 +0800
Subject: [PATCH] :bug: fix bug in recvfrom

---
 module/libsock.nas  |  3 ++-
 module/nasocket.cpp |  1 +
 std/udp.nas         | 49 +++++++++++++++++++++++++++++++++++++++++++++
 test/udptest.nas    | 18 +++++++++++++++++
 4 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 std/udp.nas
 create mode 100644 test/udptest.nas

diff --git a/module/libsock.nas b/module/libsock.nas
index fa6b356..90335f4 100644
--- a/module/libsock.nas
+++ b/module/libsock.nas
@@ -76,8 +76,9 @@ var socket=func(){
         MSG_OOB:0x1,
         MSG_PEEK:0x2,
         MSG_DONTROUTE:0x4,
+        MSG_DONTWAIT:0x40,
 
-        socket:func(af,type,proto){
+        socket:func(af,type,proto = 0){
             return invoke_iii(sock,af,type,proto);
         },
         closesocket:func(sd){
diff --git a/module/nasocket.cpp b/module/nasocket.cpp
index 7b692c6..d2a9600 100644
--- a/module/nasocket.cpp
+++ b/module/nasocket.cpp
@@ -231,6 +231,7 @@ var nas_recvfrom(var* args, usize size, gc* ngc) {
     hash["str"] = ngc->newstr(buf);
     delete[] buf;
     hash["fromip"] = ngc->newstr(inet_ntoa(addr.sin_addr));
+    hash["port"] = var::num(ntohs(addr.sin_port));
     ngc->temp = nil;
     return res;
 }
diff --git a/std/udp.nas b/std/udp.nas
new file mode 100644
index 0000000..7918a3a
--- /dev/null
+++ b/std/udp.nas
@@ -0,0 +1,49 @@
+use module.libsock;
+
+var udp_server = func(hostname, port, retry_delay = 5) {
+    var socket = libsock.socket;
+    var server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM);
+    while(socket.bind(server, hostname, port) < 0){
+        println("[", os.time(), "] failed to bind socket "~server~" at ", hostname, ":", port, ".");
+        unix.sleep(retry_delay);
+        println("[", os.time(), "] retrying...");
+    }
+
+    return {
+        sendto: func(ip, port, message) {
+            var res = socket.sendto(server, ip, port, message);
+            println("[", os.time(), "] send message to ", ip, ":", port, " ", res, " byte(s)");
+            return res;
+        },
+        recvfrom: func(length = 1024) {
+            var message = socket.recvfrom(server, length);
+            println("[", os.time(), "] get message \"", message.str, "\" from ", message.fromip, ":", message.port);
+            return message;
+        }
+    }
+}
+
+var udp_client = func(hostname = "", port = -1, retry_delay = 5) {
+    var socket = libsock.socket;
+    var client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM);
+    if (port > 0 and size(hostname) != 0) {
+        while(socket.bind(client, hostname, port)<0){
+            println("[",os.time(),"] failed to bind socket "~client~" at ", hostname, ":", port, ".");
+            unix.sleep(retry_delay);
+            println("[",os.time(),"] retrying...");
+        }
+    }
+
+    return {
+        sendto: func(ip, port, message) {
+            var res = socket.sendto(client, ip, port, message);
+            println("[", os.time(), "] send message to ", ip, ":", port, " ", res, " byte(s)");
+            return res;
+        },
+        recvfrom: func(length = 1024) {
+            var message = socket.recvfrom(client, length);
+            println("[", os.time(), "] get message \"", message.str, "\" from ", message.fromip, ":", message.port);
+            return message;
+        }
+    }
+}
\ No newline at end of file
diff --git a/test/udptest.nas b/test/udptest.nas
new file mode 100644
index 0000000..2275f93
--- /dev/null
+++ b/test/udptest.nas
@@ -0,0 +1,18 @@
+use std.udp;
+
+var argument = arg[0];
+if (argument=="server") {
+    var server = udp.udp_server("127.0.0.1", 5506);
+    while(1) {
+        var message = server.recvfrom();
+        server.sendto(message.fromip, message.port, "from server");
+    }
+} elsif (argument=="client") {
+    var client = udp.udp_client();
+    while(1) {
+        client.sendto("127.0.0.1", 5506, "hello");
+        var message = client.recvfrom();
+        println(message);
+        unix.sleep(1);
+    }
+}