improve socket recv

3 files changed, 49 insertions, 34 deletions

diff --git a/core/src/xmake/io/socket_recv.c b/core/src/xmake/io/socket_recv.c
index 0b4b680fe..ce87cffca 100644
--- a/core/src/xmake/io/socket_recv.c
+++ b/core/src/xmake/io/socket_recv.c
@@ -76,10 +76,5 @@ tb_int_t xm_io_socket_recv(lua_State* lua)
     // recv data
     tb_long_t real = tb_socket_recv(sock, data, size);
     lua_pushinteger(lua, (tb_int_t)real);
-    if (real > 0)
-    {
-        lua_pushlstring(lua, (tb_char_t const*)data, real);
-        return 2;
-    }
     return 1;
 }
diff --git a/core/src/xmake/io/socket_recvfrom.c b/core/src/xmake/io/socket_recvfrom.c
index 9bf08c542..7e192f60d 100644
--- a/core/src/xmake/io/socket_recvfrom.c
+++ b/core/src/xmake/io/socket_recvfrom.c
@@ -52,27 +52,25 @@ tb_int_t xm_io_socket_recvfrom(lua_State* lua)
     tb_socket_ref_t sock = (tb_socket_ref_t)lua_touserdata(lua, 1);
     tb_check_return_val(sock, 0);
 
-    // get data and size
-    tb_byte_t  buffer[8192];
-    tb_byte_t* data = buffer;
-    tb_long_t  size = 0;
-    if (lua_isnumber(lua, 2)) size = (tb_long_t)lua_tonumber(lua, 2);
-    if (size < 0)
+    // get data
+    tb_byte_t* data = tb_null;
+    if (lua_isnumber(lua, 2))
+        data = (tb_byte_t*)(tb_size_t)(tb_long_t)lua_tonumber(lua, 2);
+    if (!data)
     {
         lua_pushinteger(lua, -1);
-        lua_pushfstring(lua, "invalid size(%ld)!", size);
+        lua_pushfstring(lua, "invalid data(%p)!", data);
         return 2;
     }
-    if (!size) size = sizeof(data);
-    else if (size > sizeof(data)) 
+
+    // get size
+    tb_long_t size = 0;
+    if (lua_isnumber(lua, 3)) size = (tb_long_t)lua_tonumber(lua, 3);
+    if (size <= 0)
     {
-        data = tb_malloc_bytes(size);
-        if (!data)
-        {
-            lua_pushinteger(lua, -1);
-            lua_pushfstring(lua, "malloc(%ld) failed!", size);
-            return 2;
-        }
+        lua_pushinteger(lua, -1);
+        lua_pushfstring(lua, "invalid size(%ld)!", size);
+        return 2;
     }
 
     // recv data
@@ -84,10 +82,11 @@ tb_int_t xm_io_socket_recvfrom(lua_State* lua)
     if (real > 0)
     {
         retn = 2;
-        lua_pushlstring(lua, (tb_char_t const*)data, real);
+        lua_pushnil(lua);
         if (!tb_ipaddr_is_empty(&ipaddr))
         {
-            tb_char_t const* ipstr = tb_ipaddr_ip_cstr(&ipaddr, (tb_char_t*)buffer, sizeof(buffer));
+            tb_char_t buffer[256];
+            tb_char_t const* ipstr = tb_ipaddr_ip_cstr(&ipaddr, buffer, sizeof(buffer));
             if (ipstr) 
             {
                 lua_pushstring(lua, ipstr);
@@ -96,6 +95,5 @@ tb_int_t xm_io_socket_recvfrom(lua_State* lua)
             }
         }
     }
-    if (data != buffer) tb_free(data);
     return retn;
 }
diff --git a/xmake/core/base/socket.lua b/xmake/core/base/socket.lua
index a01dde5dd..455b42b0a 100644
--- a/xmake/core/base/socket.lua
+++ b/xmake/core/base/socket.lua
@@ -48,7 +48,7 @@ socket.EV_ACPT = socket.EV_RECV
 function _instance.new(socktype, family, sock)
     local instance   = table.inherit(_instance)
     instance._SOCK   = sock
-    instance._TYPE   = socktype 
+    instance._TYPE   = socktype
     instance._FAMILY = family
     setmetatable(instance, _instance)
     return instance
@@ -322,15 +322,16 @@ function _instance:recv(size, opt)
     local real = 0
     local wait = false
     local data_or_errors = nil
-    local data = bytes(8192)
     if opt.block then
         local results = {}
         while recv < size do
-            real, data_or_errors = io.socket_recv(self:csock(), data:caddr(), math.min(8192, size - recv))
+            local buff = self:_recvbuff()
+            real, data_or_errors = io.socket_recv(self:csock(), buff:caddr(), math.min(buff:size(), size - recv))
             if real > 0 then
                 recv = recv + real
                 wait = false
-                table.insert(results, bytes(data, 1, real))
+                table.insert(results, bytes(buff, 1, real))
+                self:_recvbuff_clear()
             elseif real == 0 and not wait then
                 local events, waiterrs = self:wait(socket.EV_RECV, opt.timeout or -1)
                 if events == socket.EV_RECV then
@@ -349,9 +350,11 @@ function _instance:recv(size, opt)
             recv = -1
         end
     else
-        recv, data_or_errors = io.socket_recv(self:csock(), data:caddr(), math.min(8192, size))
+        local buff = self:_recvbuff()
+        recv, data_or_errors = io.socket_recv(self:csock(), buff:caddr(), math.min(buff:size(), size))
         if recv > 0 then
-            data_or_errors = bytes(data, 1, recv)
+            data_or_errors = bytes(buff, 1, recv)
+            self:_recvbuff_clear()
         end
     end
     if recv < 0 and data_or_errors then
@@ -441,9 +444,11 @@ function _instance:recvfrom(size, opt)
     local data_or_errors = nil
     if opt.block then
         while true do
-            recv, data_or_errors, addr, port = io.socket_recvfrom(self:csock(), size)
+            local buff = self:_recvbuff()
+            recv, data_or_errors, addr, port = io.socket_recvfrom(self:csock(), buff:caddr(), math.min(buff:size(), size))
             if recv > 0 then
-                data_or_errors = bytes(data_or_errors)
+                data_or_errors = bytes(buff, 1, recv)
+                self:_recvbuff_clear()
                 break
             elseif recv == 0 and not wait then
                 local events, waiterrs = self:wait(socket.EV_RECV, opt.timeout or -1)
@@ -459,9 +464,11 @@ function _instance:recvfrom(size, opt)
             end
         end
     else
-        recv, data_or_errors, addr, port = io.socket_recvfrom(self:csock(), size)
+        local buff = self:_recvbuff()
+        recv, data_or_errors, addr, port = io.socket_recvfrom(self:csock(), buff:caddr(), math.min(buff:size(), size))
         if recv > 0 then
-            data_or_errors = bytes(data_or_errors)
+            data_or_errors = bytes(buff, 1, recv)
+            self:_recvbuff_clear()
         end
     end
     if recv < 0 and data_or_errors then
@@ -518,6 +525,21 @@ function _instance:close()
     return ok
 end
 
+-- get the recv buffer
+function _instance:_recvbuff()
+    local recvbuff = self._RECVBUFF
+    if not recvbuff then
+        recvbuff = bytes(8192)
+        self._RECVBUFF = recvbuff
+    end
+    return recvbuff
+end
+
+-- clear the recv buffer 
+function _instance:_recvbuff_clear()
+    self._RECVBUFF = nil
+end
+
 -- ensure the socket is opened
 function _instance:_ensure_opened()
     if not self:csock() then