lihao/skynet
1
0
Fork 0
mirror of https://github.com/cloudwu/skynet.git synced 2026-07-23 19:43:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

add socket.limit for defence

Browse Source
This commit is contained in:
Cloud Wu
2014-07-12 20:30:02 +08:00
parent 54e6d03d36
commit 3a5de32ad0
3 changed files with 23 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
lualib-src/lua-socket.c
View File

@@ -14,6 +14,7 @@
#define BACKLOG 32
// 2 ** 12 == 4096
#define LARGE_PAGE_NODE 12
#define BUFFER_LIMIT (256 * 1024)
struct buffer_node {
char * msg;
@@ -22,6 +23,7 @@ struct buffer_node {
};
struct socket_buffer {
int limit;
int size;
int offset;
struct buffer_node *head;
@@ -64,6 +66,7 @@ lnewpool(lua_State *L, int sz) {
static int
lnewbuffer(lua_State *L) {
struct socket_buffer * sb = lua_newuserdata(L, sizeof(*sb));
sb->limit = luaL_optint(L,1,BUFFER_LIMIT);
sb->size = 0;
sb->offset = 0;
sb->head = NULL;
@@ -126,6 +129,9 @@ lpushbuffer(lua_State *L) {
sb->size += sz;
lua_pushinteger(L, sb->size);
if (sb->limit > 0 && sb->size > sb->limit) {
return luaL_error(L, "buffer overflow (limit = %d, size = %d)", sb->limit, sb->size);
}
return 1;
}

15
lualib/socket.lua
View File

@@ -2,6 +2,7 @@ local driver = require "socketdriver"
local skynet = require "skynet"
local assert = assert
local buffer_limit = -1
local socket = {} -- api
local buffer_pool = {} -- store all message buffer object
local socket_pool = setmetatable( -- store all socket object
@@ -47,7 +48,13 @@ socket_message[1] = function(id, size, data)
return
end
local sz = driver.push(s.buffer, buffer_pool, data, size)
local ok , sz = pcall(driver.push, s.buffer, buffer_pool, data, size)
if not ok then
skynet.error("socket: error on ", id , sz)
driver.clear(s.buffer,buffer_pool)
driver.close(id)
return
end
local rr = s.read_required
local rrt = type(rr)
if rrt == "number" then
@@ -123,7 +130,7 @@ skynet.register_protocol {
local function connect(id, func)
local newbuffer
if func == nil then
newbuffer = driver.buffer()
newbuffer = driver.buffer(buffer_limit)
end
local s = {
id = id,
@@ -318,4 +325,8 @@ function socket.abandon(id)
socket_pool[id] = nil
end
function socket.limit(limit)
buffer_limit = limit
end
return socket

5
test/testsocket.lua
View File

@@ -21,6 +21,9 @@ if mode == "agent" then
id = tonumber(id)
skynet.start(function()
-- A small limit, if socket buffer overflow, close the client
socket.limit(64)
skynet.fork(function()
echo(id)
skynet.exit()
@@ -30,7 +33,7 @@ else
local function accept(id)
socket.start(id)
socket.write(id, "Hello Skynet\n")
skynet.newservice("testsocket", "agent", id)
skynet.newservice(SERVICE_NAME, "agent", id)
-- notice: Some data on this connection(id) may lost before new service start.
-- So, be careful when you want to use start / abandon / start .
socket.abandon(id)