From aad2a1b446ef81167d1a7d3850359529e6182205 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=91=E9=A3=8E?= <cloudwu@gmail.com>
Date: Fri, 17 Aug 2012 14:46:03 +0800
Subject: [PATCH] support 2-bytes length block

---
 connection/lua-socket.c | 38 ++++++++++++++++++++++++++++++++++++++
 lualib/socket.lua       |  4 ++++
 2 files changed, 42 insertions(+)

diff --git a/connection/lua-socket.c b/connection/lua-socket.c
index c8013d33..3c92e52d 100644
--- a/connection/lua-socket.c
+++ b/connection/lua-socket.c
@@ -173,6 +173,43 @@ _readline(lua_State *L) {
 	return 0;
 }
 
+/*
+	userdata buffer
+	function (msg, sz, ...)
+	...
+ */
+static int
+_readblock(lua_State *L) {
+	luaL_checktype(L,1,LUA_TUSERDATA);
+	struct buffer * buffer = lua_touserdata(L,1);
+	int top = lua_gettop(L);
+	int size = buffer->size - buffer->read;
+	if (size < 2) {
+		return 0;
+	}
+	uint8_t * buf = (uint8_t *)buffer->buffer + buffer->read;
+	uint16_t len = buf[0] << 8 | buf[1];
+	if (size < 2 + len) {
+		return 0;
+	}
+
+	if (top == 2) {
+		lua_pushlightuserdata(L, buffer->buffer + buffer->read + 2);
+		lua_pushinteger(L, len);
+		lua_call(L,2,LUA_MULTRET);
+	} else {
+		lua_pushlightuserdata(L, buffer->buffer + buffer->read + 2);
+		lua_insert(L,3);
+		lua_pushinteger(L, len);
+		lua_insert(L,4);
+		lua_call(L,top,LUA_MULTRET);
+	}
+
+	buffer->read += size + 2;
+
+	return lua_gettop(L) - 1;
+}
+
 int
 luaopen_socket_c(lua_State *L) {
 	luaL_Reg l[] = {
@@ -183,6 +220,7 @@ luaopen_socket_c(lua_State *L) {
 		{ "push", _push },
 		{ "read", _read },
 		{ "readline", _readline },
+		{ "readblock", _readblock },
 		{ NULL, NULL },
 	};
 	luaL_checkversion(L);
diff --git a/lualib/socket.lua b/lualib/socket.lua
index b64e82ed..ebb5266e 100644
--- a/lualib/socket.lua
+++ b/lualib/socket.lua
@@ -30,6 +30,10 @@ function socket.readline(sep)
 	return c.readline(object, sep)
 end
 
+function socket.readblock(...)
+	return c.readblock(object,...)
+end
+
 function socket.write(...)
 	c.write(fd, ...)
 end