1 files changed, 78 insertions, 33 deletions

diff --git a/src/udp.c b/src/udp.c
index 17d932a..712ad50 100644..100755
--- a/src/udp.c
+++ b/src/udp.c
@@ -2,12 +2,7 @@
 * UDP object
 * LuaSocket toolkit
 \*=========================================================================*/
-#include <string.h>
-#include <stdlib.h>
-
-#include "lua.h"
-#include "lauxlib.h"
-#include "compat.h"
+#include "luasocket.h"
 
 #include "auxiliar.h"
 #include "socket.h"
@@ -15,6 +10,9 @@
 #include "options.h"
 #include "udp.h"
 
+#include <string.h>
+#include <stdlib.h>
+
 /* min and max macros */
 #ifndef MIN
 #define MIN(x, y) ((x) < (y) ? x : y)
@@ -36,6 +34,7 @@ static int meth_receivefrom(lua_State *L);
 static int meth_getfamily(lua_State *L);
 static int meth_getsockname(lua_State *L);
 static int meth_getpeername(lua_State *L);
+static int meth_gettimeout(lua_State *L);
 static int meth_setsockname(lua_State *L);
 static int meth_setpeername(lua_State *L);
 static int meth_close(lua_State *L);
@@ -66,6 +65,7 @@ static luaL_Reg udp_methods[] = {
     {"setpeername", meth_setpeername},
     {"setsockname", meth_setsockname},
     {"settimeout",  meth_settimeout},
+    {"gettimeout",  meth_gettimeout},
     {NULL,          NULL}
 };
 
@@ -86,6 +86,8 @@ static t_opt optset[] = {
     {"ipv6-add-membership",  opt_set_ip6_add_membership},
     {"ipv6-drop-membership", opt_set_ip6_drop_membersip},
     {"ipv6-v6only",          opt_set_ip6_v6only},
+        {"recv-buffer-size",     opt_set_recv_buf_size},
+        {"send-buffer-size",     opt_set_send_buf_size},
     {NULL,                   NULL}
 };
 
@@ -102,6 +104,8 @@ static t_opt optget[] = {
     {"ipv6-multicast-hops",  opt_get_ip6_unicast_hops},
     {"ipv6-multicast-loop",  opt_get_ip6_multicast_loop},
     {"ipv6-v6only",          opt_get_ip6_v6only},
+        {"recv-buffer-size",     opt_get_recv_buf_size},
+        {"send-buffer-size",     opt_get_send_buf_size},
     {NULL,                   NULL}
 };
 
@@ -116,8 +120,7 @@ static luaL_Reg func[] = {
 /*-------------------------------------------------------------------------*\
 * Initializes module
 \*-------------------------------------------------------------------------*/
-int udp_open(lua_State *L)
-{
+int udp_open(lua_State *L) {
     /* create classes */
     auxiliar_newclass(L, "udp{connected}", udp_methods);
     auxiliar_newclass(L, "udp{unconnected}", udp_methods);
@@ -128,6 +131,10 @@ int udp_open(lua_State *L)
     auxiliar_add2group(L, "udp{unconnected}", "select{able}");
     /* define library functions */
     luaL_setfuncs(L, func, 0);
+    /* export default UDP size */
+    lua_pushliteral(L, "_DATAGRAMSIZE");
+    lua_pushinteger(L, UDP_DATAGRAMSIZE);
+    lua_rawset(L, -3);
     return 0;
 }
 
@@ -177,13 +184,37 @@ static int meth_sendto(lua_State *L) {
     memset(&aihint, 0, sizeof(aihint));
     aihint.ai_family = udp->family;
     aihint.ai_socktype = SOCK_DGRAM;
-    aihint.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV;
+    aihint.ai_flags = AI_NUMERICHOST;
+#ifdef AI_NUMERICSERV
+    aihint.ai_flags |= AI_NUMERICSERV;
+#endif
     err = getaddrinfo(ip, port, &aihint, &ai);
         if (err) {
         lua_pushnil(L);
-        lua_pushstring(L, gai_strerror(err));
+        lua_pushstring(L, LUA_GAI_STRERROR(err));
         return 2;
     }
+
+    /* create socket if on first sendto if AF_UNSPEC was set */
+    if (udp->family == AF_UNSPEC && udp->sock == SOCKET_INVALID) {
+        struct addrinfo *ap;
+        const char *errstr = NULL;
+        for (ap = ai; ap != NULL; ap = ap->ai_next) {
+            errstr = inet_trycreate(&udp->sock, ap->ai_family, SOCK_DGRAM, 0);
+            if (errstr == NULL) {
+                socket_setnonblocking(&udp->sock);
+                udp->family = ap->ai_family;
+                break;
+            }
+        }
+        if (errstr != NULL) {
+            lua_pushnil(L);
+            lua_pushstring(L, errstr);
+            freeaddrinfo(ai);
+            return 2;
+        }
+    }
+
     timeout_markstart(tm);
     err = socket_sendto(&udp->sock, data, count, &sent, ai->ai_addr,
         (socklen_t) ai->ai_addrlen, tm);
@@ -202,69 +233,78 @@ static int meth_sendto(lua_State *L) {
 \*-------------------------------------------------------------------------*/
 static int meth_receive(lua_State *L) {
     p_udp udp = (p_udp) auxiliar_checkgroup(L, "udp{any}", 1);
-    char buffer[UDP_DATAGRAMSIZE];
-    size_t got, count = (size_t) luaL_optnumber(L, 2, sizeof(buffer));
+    char buf[UDP_DATAGRAMSIZE];
+    size_t got, wanted = (size_t) luaL_optnumber(L, 2, sizeof(buf));
+    char *dgram = wanted > sizeof(buf)? (char *) malloc(wanted): buf;
     int err;
     p_timeout tm = &udp->tm;
-    count = MIN(count, sizeof(buffer));
     timeout_markstart(tm);
-    err = socket_recv(&udp->sock, buffer, count, &got, tm);
+    if (!dgram) {
+        lua_pushnil(L);
+        lua_pushliteral(L, "out of memory");
+        return 2;
+    }
+    err = socket_recv(&udp->sock, dgram, wanted, &got, tm);
     /* Unlike TCP, recv() of zero is not closed, but a zero-length packet. */
-    if (err == IO_CLOSED)
-        err = IO_DONE;
-    if (err != IO_DONE) {
+    if (err != IO_DONE && err != IO_CLOSED) {
         lua_pushnil(L);
         lua_pushstring(L, udp_strerror(err));
+        if (wanted > sizeof(buf)) free(dgram);
         return 2;
     }
-    lua_pushlstring(L, buffer, got);
+    lua_pushlstring(L, dgram, got);
+    if (wanted > sizeof(buf)) free(dgram);
     return 1;
 }
 
 /*-------------------------------------------------------------------------*\
 * Receives data and sender from a UDP socket
 \*-------------------------------------------------------------------------*/
-static int meth_receivefrom(lua_State *L)
-{
+static int meth_receivefrom(lua_State *L) {
     p_udp udp = (p_udp) auxiliar_checkclass(L, "udp{unconnected}", 1);
-    char buffer[UDP_DATAGRAMSIZE];
-    size_t got, count = (size_t) luaL_optnumber(L, 2, sizeof(buffer));
-    int err;
-    p_timeout tm = &udp->tm;
+    char buf[UDP_DATAGRAMSIZE];
+    size_t got, wanted = (size_t) luaL_optnumber(L, 2, sizeof(buf));
+    char *dgram = wanted > sizeof(buf)? (char *) malloc(wanted): buf;
     struct sockaddr_storage addr;
     socklen_t addr_len = sizeof(addr);
     char addrstr[INET6_ADDRSTRLEN];
     char portstr[6];
+    int err;
+    p_timeout tm = &udp->tm;
     timeout_markstart(tm);
-    count = MIN(count, sizeof(buffer));
-    err = socket_recvfrom(&udp->sock, buffer, count, &got, (SA *) &addr,
+    if (!dgram) {
+        lua_pushnil(L);
+        lua_pushliteral(L, "out of memory");
+        return 2;
+    }
+    err = socket_recvfrom(&udp->sock, dgram, wanted, &got, (SA *) &addr,
             &addr_len, tm);
     /* Unlike TCP, recv() of zero is not closed, but a zero-length packet. */
-    if (err == IO_CLOSED)
-        err = IO_DONE;
-    if (err != IO_DONE) {
+    if (err != IO_DONE && err != IO_CLOSED) {
         lua_pushnil(L);
         lua_pushstring(L, udp_strerror(err));
+        if (wanted > sizeof(buf)) free(dgram);
         return 2;
     }
     err = getnameinfo((struct sockaddr *)&addr, addr_len, addrstr,
         INET6_ADDRSTRLEN, portstr, 6, NI_NUMERICHOST | NI_NUMERICSERV);
         if (err) {
         lua_pushnil(L);
-        lua_pushstring(L, gai_strerror(err));
+        lua_pushstring(L, LUA_GAI_STRERROR(err));
+        if (wanted > sizeof(buf)) free(dgram);
         return 2;
     }
-    lua_pushlstring(L, buffer, got);
+    lua_pushlstring(L, dgram, got);
     lua_pushstring(L, addrstr);
     lua_pushinteger(L, (int) strtol(portstr, (char **) NULL, 10));
+    if (wanted > sizeof(buf)) free(dgram);
     return 3;
 }
 
 /*-------------------------------------------------------------------------*\
 * Returns family as string
 \*-------------------------------------------------------------------------*/
-static int meth_getfamily(lua_State *L)
-{
+static int meth_getfamily(lua_State *L) {
     p_udp udp = (p_udp) auxiliar_checkgroup(L, "udp{any}", 1);
     if (udp->family == AF_INET6) {
         lua_pushliteral(L, "inet6");
@@ -335,6 +375,11 @@ static int meth_settimeout(lua_State *L) {
     return timeout_meth_settimeout(L, &udp->tm);
 }
 
+static int meth_gettimeout(lua_State *L) {
+    p_udp udp = (p_udp) auxiliar_checkgroup(L, "udp{any}", 1);
+    return timeout_meth_gettimeout(L, &udp->tm);
+}
+
 /*-------------------------------------------------------------------------*\
 * Turns a master udp object into a client object.
 \*-------------------------------------------------------------------------*/