socket: accept data object in send/receive to reuse buffers

sock:send(data [, addr])   -- sends data->ptr/size, no string alloc
sock:receive(data [, flags [, from]])  -- receives into data->ptr, returns bytes

String path is preserved for backward compatibility. luadata_t is moved
to luadata.h and luadata_arg() added for cross-module use.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: lneto

lib/luadata.c

@@ -16,13 +16,6 @@
 
 #include "luadata.h"
 
-
-typedef struct luadata_s {
-	void *ptr;
-	size_t size;
-	uint8_t opt;
-} luadata_t;
-
 #define LUADATA_NUMBER_SZ	(sizeof(lua_Integer))
 
 static int luadata_lnew(lua_State *L);

lib/luadata.h

@@ -26,5 +26,12 @@ static inline void luadata_close(lunatik_object_t *object)
 
 #define luadata_attach(L, obj, field, opt)	lunatik_attach(L, obj, field, luadata_new, opt)
 
+typedef struct luadata_s {
+	void *ptr;
+	size_t size;
+	uint8_t opt;
+} luadata_t;
+
+
 #endif
 

lib/luasocket.c

@@ -32,6 +32,7 @@
 #endif
 
 #include <lunatik.h>
+#include "luadata.h"
 
 #define luasocket_msgaddr(msg, addr, size)	\
 do {						\
@@ -117,7 +118,7 @@ LUNATIK_PRIVATECHECKER(luasocket_check, struct socket *);
 * address family) specify the destination.
 *
 * @function send
-* @tparam string message The string message to send.
+* @tparam string|data message The message to send; a data object avoids string allocation.
 * @tparam[opt] integer|string addr The destination address.
 *
 * - For `AF_INET` (IPv4) sockets: An integer representing the IPv4 address (e.g., from `net.aton()`).
@@ -146,7 +147,14 @@ static int luasocket_send(lua_State *L)
 
 	luasocket_setmsg(msg);
 
-	vec.iov_base = (void *)luaL_checklstring(L, 2, &len);
+	if (lua_isuserdata(L, 2)) {
+		luadata_t *data = lunatik_checkclass(L, 2, "data", LUNATIK_OPT_SINGLE)->private;
+		lunatik_argchecknull(L, data->ptr, 2);
+		vec.iov_base = data->ptr;
+		len = data->size;
+	}
+	else
+		vec.iov_base = (void *)luaL_checklstring(L, 2, &len);
 	vec.iov_len = len;
 
 	if (unlikely(nargs >= 3)) {
@@ -163,32 +171,21 @@ static int luasocket_send(lua_State *L)
 * Receives a message from the socket.
 *
 * @function receive
-* @tparam integer length The maximum number of bytes to receive.
-* @tparam[opt=0] integer flags Optional message flags (e.g., `socket.msg.PEEK`).
-*   See the `socket.msg` table for available flags. These can be OR'd together.
-* @tparam[opt=false] boolean from If `true`, the function also returns the sender's address
-*   and port (for `AF_INET`). This is typically used with connectionless sockets (`SOCK_DGRAM`).
-* @treturn string The received message (as a string of bytes).
-* @treturn[opt] integer|string addr If `from` is true, the sender's address.
-*   - For `AF_INET`: An integer representing the IPv4 address (can be converted with `net.ntoa()`).
-*   - For other families: A packed string representing the sender's address.
-* @treturn[opt] integer port If `from` is true and the family is `AF_INET`, the sender's port number.
+* @tparam integer|data buffer Maximum bytes to receive, or a data object to receive into.
+* @tparam[opt=0] integer flags Message flags (e.g., `socket.msg.PEEK`).
+* @tparam[opt=false] boolean from If `true`, also returns sender address and port (`AF_INET`).
+* @treturn string|integer Received bytes as a string, or byte count when a data object is passed.
+* @treturn[opt] integer|string addr Sender address if `from` is true.
+* @treturn[opt] integer port Sender port if `from` is true and family is `AF_INET`.
 * @raise Error if the receive operation fails.
-* @usage
-*   -- For a connected TCP socket:
-*   local data = tcp_conn_sock:receive(1024)
-*   if data then print("Received:", data) end
-*
-*   -- For a UDP socket, getting sender info:
-*   local data, sender_ip_int, sender_port = udp_sock:receive(1500, 0, true)
-*   if data then print("Received from " .. net.ntoa(sender_ip_int) .. ":" .. sender_port .. ": " .. data) end
 * @see socket.msg
 * @see net.ntoa
 */
 static int luasocket_receive(lua_State *L)
 {
 	struct socket *socket = luasocket_check(L, 1);
-	size_t len = (size_t)luaL_checkinteger(L, 2);
+	int isdata = lua_isuserdata(L, 2);
+	size_t len;
 	luaL_Buffer B;
 	struct kvec vec;
 	struct msghdr msg;
@@ -199,14 +196,23 @@ static int luasocket_receive(lua_State *L)
 
 	luasocket_setmsg(msg);
 
-	vec.iov_base = (void *)luaL_buffinitsize(L, &B, len);
+	if (isdata) {
+		luadata_t *data = lunatik_checkclass(L, 2, "data", LUNATIK_OPT_SINGLE)->private;
+		lunatik_argchecknull(L, data->ptr, 2);
+		vec.iov_base = data->ptr;
+		len = data->size;
+	}
+	else {
+		len = (size_t)luaL_checkinteger(L, 2);
+		vec.iov_base = (void *)luaL_buffinitsize(L, &B, len);
+	}
 	vec.iov_len = len;
 
 	if (unlikely(from))
 		luasocket_msgaddr(msg, addr, sizeof(addr));
 
 	lunatik_tryret(L, ret, kernel_recvmsg, socket, &msg, &vec, 1, len, flags);
-	luaL_pushresultsize(&B, ret);
+	isdata ? lua_pushinteger(L, ret) : luaL_pushresultsize(&B, ret);
 
 	return unlikely(from) ? luasocket_pushaddr(L, (struct sockaddr_storage *)msg.msg_name) + 1 : 1;
 }

lunatik.h

@@ -9,6 +9,7 @@
 #include <linux/mutex.h>
 #include <linux/spinlock.h>
 #include <linux/slab.h>
+#include <linux/string.h>
 #include <linux/kref.h>
 #include <linux/version.h>
 
@@ -194,7 +195,7 @@ static inline lunatik_object_t *lunatik_checkruntime(lua_State *L, lunatik_opt_t
 #define lunatik_setruntime(L, libname, priv)	((priv)->runtime = lunatik_checkruntime((L), lua##libname##_class.opt))
 #define lunatik_monitormt(class, monitor)	((monitor) ? (void *)&(class)->opt : (void *)(class))
 
-static inline void lunatik_checkclass(lua_State *L, const lunatik_class_t *class)
+static inline void lunatik_checkcontext(lua_State *L, const lunatik_class_t *class)
 {
 	if (lunatik_cannotsleep(L, !lunatik_issoftirq(class->opt)))
 		luaL_error(L, "'%s': %s", class->name, LUNATIK_ERR_CONTEXT);
@@ -292,6 +293,15 @@ static inline lunatik_object_t **lunatik_testobject(lua_State *L, int ix)
 	return (pobject && lunatik_isobject(L, ix, *pobject)) ? pobject : NULL;
 }
 
+static inline lunatik_object_t *lunatik_checkclass(lua_State *L, int ix,
+	const char *name, lunatik_opt_t opt)
+{
+	lunatik_object_t *object = lunatik_checkobject(L, ix);
+	luaL_argcheck(L, strcmp(object->class->name, name) == 0, ix, name);
+	luaL_argcheck(L, !opt || (object->opt & opt) == opt, ix, name);
+	return object;
+}
+
 static inline lunatik_object_t **lunatik_checkpobject(lua_State *L, int ix)
 {
 	lunatik_object_t **pobject = lunatik_testobject(L, ix);
@@ -320,7 +330,7 @@ int luaopen_##libname(lua_State *L)						\
 	const lunatik_namespace_t *nss = namespaces; /* avoid -Waddress */	\
 	luaL_newlib(L, funcs);							\
 	if (cls) {								\
-		lunatik_checkclass(L, cls);					\
+		lunatik_checkcontext(L, cls);					\
 		if (lunatik_ismonitor(cls->opt))				\
 			lunatik_newclass(L, cls, true);				\
 		lunatik_newclass(L, cls, false);				\

lunatik_obj.c

@@ -22,7 +22,7 @@ lunatik_object_t *lunatik_newobject(lua_State *L, const lunatik_class_t *class,
 	lunatik_object_t *object = lunatik_checkalloc(L, sizeof(lunatik_object_t));
 
 	/* SOFTIRQ runtime requires a SOFTIRQ class */
-	lunatik_checkclass(L, class);
+	lunatik_checkcontext(L, class);
 
 	lunatik_setobject(object, class, opt);
 	lunatik_setclass(L, class, lunatik_ismonitor(object->opt));
@@ -62,7 +62,7 @@ void lunatik_cloneobject(lua_State *L, lunatik_object_t *object)
 	lunatik_require(L, class->name);
 	lunatik_object_t **pobject = lunatik_newpobject(L, 1);
 
-	lunatik_checkclass(L, class);
+	lunatik_checkcontext(L, class);
 	lunatik_setclass(L, class, lunatik_ismonitor(object->opt));
 	*pobject = object;
 }