Merge branch 'udp' into dev

lualib-src/lua-socket.c

@@ -9,6 +9,8 @@
 #include <lua.h>
 #include <lauxlib.h>
 
+#include <arpa/inet.h>
+
 #include "skynet_socket.h"
 
 #define BACKLOG 32
@@ -372,28 +374,64 @@ lunpack(lua_State *L) {
 	} else {
 		lua_pushlightuserdata(L, message->buffer);
 	}
+	if (message->type == SKYNET_SOCKET_TYPE_UDP) {
+		int addrsz = 0;
+		const char * addrstring = skynet_socket_udp_address(message, &addrsz);
+		if (addrstring) {
+			lua_pushlstring(L, addrstring, addrsz);
+			return 5;
+		}
+	}
 	return 4;
 }
 
+static const char *
+address_port(lua_State *L, char *tmp, const char * addr, int port_index, int *port) {
+	const char * host;
+	if (lua_isnoneornil(L,port_index)) {
+		host = strchr(addr, '[');
+		if (host) {
+			// is ipv6
+			++host;
+			const char * sep = strchr(addr,']');
+			if (sep == NULL) {
+				luaL_error(L, "Invalid address %s.",addr);
+			}
+			memcpy(tmp, host, sep-host);
+			tmp[sep-host] = '\0';
+			host = tmp;
+			sep = strchr(sep + 1, ':');
+			if (sep == NULL) {
+				luaL_error(L, "Invalid address %s.",addr);
+			}
+			*port = strtoul(sep+1,NULL,10);
+		} else {
+			// is ipv4
+			const char * sep = strchr(addr,':');
+			if (sep == NULL) {
+				luaL_error(L, "Invalid address %s.",addr);
+			}
+			memcpy(tmp, addr, sep-addr);
+			tmp[sep-addr] = '\0';
+			host = tmp;
+			*port = strtoul(sep+1,NULL,10);
+		}
+	} else {
+		host = addr;
+		*port = luaL_optinteger(L,port_index, 0);
+	}
+	return host;
+}
+
 static int
 lconnect(lua_State *L) {
 	size_t sz = 0;
 	const char * addr = luaL_checklstring(L,1,&sz);
 	char tmp[sz];
-	int port;
-	const char * host;
-	if (lua_isnoneornil(L,2)) {
-		const char * sep = strchr(addr,':');
-		if (sep == NULL) {
-			return luaL_error(L, "Connect to invalid address %s.",addr);
-		}
-		memcpy(tmp, addr, sep-addr);
-		tmp[sep-addr] = '\0';
-		host = tmp;
-		port = strtoul(sep+1,NULL,10);
-	} else {
-		host = addr;
-		port = luaL_checkinteger(L,2);
+	int port = 0;
+	const char * host = address_port(L, tmp, addr, 2, &port);
+	if (port == 0) {
+		return luaL_error(L, "Invalid port");
 	}
 	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
 	int id = skynet_socket_connect(ctx, host, port);
@@ -426,14 +464,14 @@ llisten(lua_State *L) {
 }
 
 static void *
-get_buffer(lua_State *L, int *sz) {
+get_buffer(lua_State *L, int index, int *sz) {
 	void *buffer;
-	if (lua_isuserdata(L,2)) {
-		buffer = lua_touserdata(L,2);
-		*sz = luaL_checkinteger(L,3);
+	if (lua_isuserdata(L,index)) {
+		buffer = lua_touserdata(L,index);
+		*sz = luaL_checkinteger(L,index+1);
 	} else {
 		size_t len = 0;
-		const char * str =  luaL_checklstring(L, 2, &len);
+		const char * str =  luaL_checklstring(L, index, &len);
 		buffer = skynet_malloc(len);
 		memcpy(buffer, str, len);
 		*sz = (int)len;
@@ -446,7 +484,7 @@ lsend(lua_State *L) {
 	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
 	int id = luaL_checkinteger(L, 1);
 	int sz = 0;
-	void *buffer = get_buffer(L, &sz);
+	void *buffer = get_buffer(L, 2, &sz);
 	int err = skynet_socket_send(ctx, id, buffer, sz);
 	lua_pushboolean(L, !err);
 	return 1;
@@ -457,7 +495,7 @@ lsendlow(lua_State *L) {
 	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
 	int id = luaL_checkinteger(L, 1);
 	int sz = 0;
-	void *buffer = get_buffer(L, &sz);
+	void *buffer = get_buffer(L, 2, &sz);
 	skynet_socket_send_lowpriority(ctx, id, buffer, sz);
 	return 0;
 }
@@ -486,6 +524,90 @@ lnodelay(lua_State *L) {
 	skynet_socket_nodelay(ctx,id);
 	return 0;
 }
+/*
+int skynet_socket_udp(struct skynet_context *ctx, const char * addr, int port);
+int skynet_socket_udp_connect(struct skynet_context *ctx, int id, const char * addr, int port);
+int skynet_socket_udp_send(struct skynet_context *ctx, int id, const char * address, const void *buffer, int sz);
+const char * skynet_socket_udp_address(struct skynet_context *ctx, struct skynet_socket_message *, int *addrsz);
+*/
+
+static int
+ludp(lua_State *L) {
+	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
+	size_t sz = 0;
+	const char * addr = lua_tolstring(L,1,&sz);
+	char tmp[sz];
+	int port = 0;
+	const char * host = NULL;
+	if (addr) {
+		host = address_port(L, tmp, addr, 2, &port);
+	}
+
+	int id = skynet_socket_udp(ctx, host, port);
+	if (id < 0) {
+		return luaL_error(L, "udp init failed");
+	}
+	lua_pushinteger(L, id);
+	return 1;
+}
+
+static int
+ludp_connect(lua_State *L) {
+	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
+	int id = luaL_checkinteger(L, 1);
+	size_t sz = 0;
+	const char * addr = luaL_checklstring(L,2,&sz);
+	char tmp[sz];
+	int port = 0;
+	const char * host = NULL;
+	if (addr) {
+		host = address_port(L, tmp, addr, 3, &port);
+	}
+
+	if (skynet_socket_udp_connect(ctx, id, host, port)) {
+		return luaL_error(L, "udp connect failed");
+	}
+
+	return 0;
+}
+
+static int
+ludp_send(lua_State *L) {
+	struct skynet_context * ctx = lua_touserdata(L, lua_upvalueindex(1));
+	int id = luaL_checkinteger(L, 1);
+	const char * address = luaL_checkstring(L, 2);
+	int sz = 0;
+	void *buffer = get_buffer(L, 3, &sz);
+	int err = skynet_socket_udp_send(ctx, id, address, buffer, sz);
+
+	lua_pushboolean(L, !err);
+
+	return 1;
+}
+
+static int
+ludp_address(lua_State *L) {
+	size_t sz = 0;
+	const uint8_t * addr = (const uint8_t *)luaL_checklstring(L, 1, &sz);
+	int port = addr[1] * 256 + addr[2];
+	const void * src = addr+3;
+	char tmp[256];
+	int family;
+	if (sz == 1+2+4) {
+		family = AF_INET;
+	} else {
+		if (sz != 1+2+16) {
+			return luaL_error(L, "Invalid udp address");
+		}
+		family = AF_INET6;
+	}
+	if (inet_ntop(family, src, tmp, sizeof(tmp)) == NULL) {
+		return luaL_error(L, "Invalid udp address");
+	}
+	lua_pushstring(L, tmp);
+	lua_pushinteger(L, port);
+	return 2;
+}
 
 int
 luaopen_socketdriver(lua_State *L) {
@@ -514,6 +636,10 @@ luaopen_socketdriver(lua_State *L) {
 		{ "bind", lbind },
 		{ "start", lstart },
 		{ "nodelay", lnodelay },
+		{ "udp", ludp },
+		{ "udp_connect", ludp_connect },
+		{ "udp_send", ludp_send },
+		{ "udp_address", ludp_address },
 		{ NULL, NULL },
 	};
 	lua_getfield(L, LUA_REGISTRYINDEX, "skynet_context");

lualib/socket.lua

@@ -119,12 +119,23 @@ socket_message[5] = function(id)
 	wakeup(s)
 end
 
+-- SKYNET_SOCKET_TYPE_UDP = 6
+socket_message[6] = function(id, size, data, address)
+	local s = socket_pool[id]
+	if s == nil or s.callback == nil then
+		skynet.error("socket: drop udp package from " .. id)
+		driver.drop(data, size)
+		return
+	end
+	s.callback(data, size, address)
+end
+
 skynet.register_protocol {
 	name = "socket",
 	id = skynet.PTYPE_SOCKET,	-- PTYPE_SOCKET = 6
 	unpack = driver.unpack,
-	dispatch = function (_, _, t, n1, n2, data)
-		socket_message[t](n1,n2,data)
+	dispatch = function (_, _, t, ...)
+		socket_message[t](...)
 	end
 }
 
@@ -137,9 +148,10 @@ local function connect(id, func)
 		id = id,
 		buffer = newbuffer,
 		connected = false,
-		read_require = false,
+		read_required = false,
 		co = false,
 		callback = func,
+		protocol = "TCP",
 	}
 	socket_pool[id] = s
 	suspend(s)
@@ -354,4 +366,39 @@ function socket.limit(id, limit)
 	s.buffer_limit = limit
 end
 
+---------------------- UDP
+
+local udp_socket = {}
+
+local function create_udp_object(id, cb)
+	socket_pool[id] = {
+		id = id,
+		connected = true,
+		protocol = "UDP",
+		callback = cb,
+	}
+end
+
+function socket.udp(callback, host, port)
+	local id = driver.udp(host, port)
+	create_udp_object(id, callback)
+	return id
+end
+
+function socket.udp_connect(id, addr, port, callback)
+	local obj = socket_pool[id]
+	if obj then
+		assert(obj.protocol == "UDP")
+		if callback then
+			obj.callback = callback
+		end
+	else
+		create_udp_object(id, callback)
+	end
+	driver.udp_connect(id, addr, port)
+end
+
+socket.sendto = assert(driver.udp_send)
+socket.udp_address = assert(driver.udp_address)
+
 return socket

skynet-src/skynet_socket.c

@@ -91,6 +91,9 @@ skynet_socket_poll() {
 	case SOCKET_ACCEPT:
 		forward_message(SKYNET_SOCKET_TYPE_ACCEPT, true, &result);
 		break;
+	case SOCKET_UDP:
+		forward_message(SKYNET_SOCKET_TYPE_UDP, false, &result);
+		break;
 	default:
 		skynet_error(NULL, "Unknown socket message type %d.",type);
 		return -1;
@@ -101,9 +104,8 @@ skynet_socket_poll() {
 	return 1;
 }
 
-int
-skynet_socket_send(struct skynet_context *ctx, int id, void *buffer, int sz) {
-	int64_t wsz = socket_server_send(SOCKET_SERVER, id, buffer, sz);
+static int
+check_wsz(struct skynet_context *ctx, int id, void *buffer, int64_t wsz) {
 	if (wsz < 0) {
 		skynet_free(buffer);
 		return -1;
@@ -116,6 +118,12 @@ skynet_socket_send(struct skynet_context *ctx, int id, void *buffer, int sz) {
 	return 0;
 }
 
+int
+skynet_socket_send(struct skynet_context *ctx, int id, void *buffer, int sz) {
+	int64_t wsz = socket_server_send(SOCKET_SERVER, id, buffer, sz);
+	return check_wsz(ctx, id, buffer, wsz);
+}
+
 void
 skynet_socket_send_lowpriority(struct skynet_context *ctx, int id, void *buffer, int sz) {
 	socket_server_send_lowpriority(SOCKET_SERVER, id, buffer, sz);
@@ -155,3 +163,33 @@ void
 skynet_socket_nodelay(struct skynet_context *ctx, int id) {
 	socket_server_nodelay(SOCKET_SERVER, id);
 }
+
+int 
+skynet_socket_udp(struct skynet_context *ctx, const char * addr, int port) {
+	uint32_t source = skynet_context_handle(ctx);
+	return socket_server_udp(SOCKET_SERVER, source, addr, port);
+}
+
+int 
+skynet_socket_udp_connect(struct skynet_context *ctx, int id, const char * addr, int port) {
+	return socket_server_udp_connect(SOCKET_SERVER, id, addr, port);
+}
+
+int 
+skynet_socket_udp_send(struct skynet_context *ctx, int id, const char * address, const void *buffer, int sz) {
+	int64_t wsz = socket_server_udp_send(SOCKET_SERVER, id, (const struct socket_udp_address *)address, buffer, sz);
+	return check_wsz(ctx, id, (void *)buffer, wsz);
+}
+
+const char *
+skynet_socket_udp_address(struct skynet_socket_message *msg, int *addrsz) {
+	if (msg->type != SKYNET_SOCKET_TYPE_UDP) {
+		return NULL;
+	}
+	struct socket_message sm;
+	sm.id = msg->id;
+	sm.opaque = 0;
+	sm.ud = msg->ud;
+	sm.data = msg->buffer;
+	return (const char *)socket_server_udp_address(SOCKET_SERVER, &sm, addrsz);
+}

skynet-src/skynet_socket.h

@@ -8,6 +8,7 @@ struct skynet_context;
 #define SKYNET_SOCKET_TYPE_CLOSE 3
 #define SKYNET_SOCKET_TYPE_ACCEPT 4
 #define SKYNET_SOCKET_TYPE_ERROR 5
+#define SKYNET_SOCKET_TYPE_UDP 6
 
 struct skynet_socket_message {
 	int type;
@@ -30,4 +31,9 @@ void skynet_socket_close(struct skynet_context *ctx, int id);
 void skynet_socket_start(struct skynet_context *ctx, int id);
 void skynet_socket_nodelay(struct skynet_context *ctx, int id);
 
+int skynet_socket_udp(struct skynet_context *ctx, const char * addr, int port);
+int skynet_socket_udp_connect(struct skynet_context *ctx, int id, const char * addr, int port);
+int skynet_socket_udp_send(struct skynet_context *ctx, int id, const char * address, const void *buffer, int sz);
+const char * skynet_socket_udp_address(struct skynet_socket_message *, int *addrsz);
+
 #endif

skynet-src/socket_server.c

@@ -37,28 +37,44 @@
 
 #define HASH_ID(id) (((unsigned)id) % MAX_SOCKET)
 
+#define PROTOCOL_TCP 0
+#define PROTOCOL_UDP 1
+#define PROTOCOL_UDPv6 2
+
+#define UDP_ADDRESS_SIZE 19	// ipv6 128bit + port 16bit + 1 byte type
+
+#define MAX_UDP_PACKAGE 65535
+
 struct write_buffer {
 	struct write_buffer * next;
-	char *ptr;
 	void *buffer;
+	char *ptr;
 	int sz;
 	bool userobject;
+	uint8_t udp_address[UDP_ADDRESS_SIZE];
 };
 
+#define SIZEOF_TCPBUFFER (offsetof(struct write_buffer, udp_address[0]))
+#define SIZEOF_UDPBUFFER (sizeof(struct write_buffer))
+
 struct wb_list {
 	struct write_buffer * head;
 	struct write_buffer * tail;
 };
 
 struct socket {
-	int fd;
-	int id;
-	int type;
-	int size;
-	int64_t wb_size;
 	uintptr_t opaque;
 	struct wb_list high;
 	struct wb_list low;
+	int64_t wb_size;
+	int fd;
+	int id;
+	uint16_t protocol;
+	uint16_t type;
+	union {
+		int size;
+		uint8_t udp_address[UDP_ADDRESS_SIZE];
+	} p;
 };
 
 struct socket_server {
@@ -73,6 +89,7 @@ struct socket_server {
 	struct event ev[MAX_EVENT];
 	struct socket slot[MAX_SOCKET];
 	char buffer[MAX_INFO];
+	uint8_t udpbuffer[MAX_UDP_PACKAGE];
 	fd_set rfds;
 };
 
@@ -89,6 +106,16 @@ struct request_send {
 	char * buffer;
 };
 
+struct request_send_udp {
+	struct request_send send;
+	uint8_t address[UDP_ADDRESS_SIZE];
+};
+
+struct request_setudp {
+	int id;
+	uint8_t address[UDP_ADDRESS_SIZE];
+};
+
 struct request_close {
 	int id;
 	uintptr_t opaque;
@@ -118,17 +145,44 @@ struct request_setopt {
 	int value;
 };
 
+struct request_udp {
+	int id;
+	int fd;
+	int family;
+	uintptr_t opaque;
+};
+
+/*
+	The first byte is TYPE
+
+	S Start socket
+	B Bind socket
+	L Listen socket
+	K Close socket
+	O Connect to (Open)
+	X Exit
+	D Send package (high)
+	P Send package (low)
+	A Send UDP package
+	T Set opt
+	U Create UDP socket
+	C set udp address
+ */
+
 struct request_package {
 	uint8_t header[8];	// 6 bytes dummy
 	union {
 		char buffer[256];
 		struct request_open open;
 		struct request_send send;
+		struct request_send_udp send_udp;
 		struct request_close close;
 		struct request_listen listen;
 		struct request_bind bind;
 		struct request_start start;
 		struct request_setopt setopt;
+		struct request_udp udp;
+		struct request_setudp set_udp;
 	} u;
 	uint8_t dummy[256];
 };
@@ -309,7 +363,7 @@ check_wb_list(struct wb_list *s) {
 }
 
 static struct socket *
-new_fd(struct socket_server *ss, int id, int fd, uintptr_t opaque, bool add) {
+new_fd(struct socket_server *ss, int id, int fd, int protocol, uintptr_t opaque, bool add) {
 	struct socket * s = &ss->slot[HASH_ID(id)];
 	assert(s->type == SOCKET_TYPE_RESERVE);
 
@@ -322,7 +376,8 @@ new_fd(struct socket_server *ss, int id, int fd, uintptr_t opaque, bool add) {
 
 	s->id = id;
 	s->fd = fd;
-	s->size = MIN_READ_BUFFER;
+	s->protocol = protocol;
+	s->p.size = MIN_READ_BUFFER;
 	s->opaque = opaque;
 	s->wb_size = 0;
 	check_wb_list(&s->high);
@@ -345,7 +400,7 @@ open_socket(struct socket_server *ss, struct request_open * request, struct sock
 	struct addrinfo *ai_ptr = NULL;
 	char port[16];
 	sprintf(port, "%d", request->port);
-	memset( &ai_hints, 0, sizeof( ai_hints ) );
+	memset(&ai_hints, 0, sizeof( ai_hints ) );
 	ai_hints.ai_family = AF_UNSPEC;
 	ai_hints.ai_socktype = SOCK_STREAM;
 	ai_hints.ai_protocol = IPPROTO_TCP;
@@ -375,7 +430,7 @@ open_socket(struct socket_server *ss, struct request_open * request, struct sock
 		goto _failed;
 	}
 
-	ns = new_fd(ss, id, sock, request->opaque, true);
+	ns = new_fd(ss, id, sock, PROTOCOL_TCP, request->opaque, true);
 	if (ns == NULL) {
 		close(sock);
 		goto _failed;
@@ -404,7 +459,7 @@ _failed:
 }
 
 static int
-send_list(struct socket_server *ss, struct socket *s, struct wb_list *list, struct socket_message *result) {
+send_list_tcp(struct socket_server *ss, struct socket *s, struct wb_list *list, struct socket_message *result) {
 	while (list->head) {
 		struct write_buffer * tmp = list->head;
 		for (;;) {
@@ -435,6 +490,74 @@ send_list(struct socket_server *ss, struct socket *s, struct wb_list *list, stru
 	return -1;
 }
 
+static socklen_t
+udp_socket_address(struct socket *s, const uint8_t udp_address[UDP_ADDRESS_SIZE], union sockaddr_all *sa) {
+	int type = (uint8_t)udp_address[0];
+	if (type != s->protocol)
+		return 0;
+	uint16_t port = 0;
+	memcpy(&port, udp_address+1, sizeof(uint16_t));
+	switch (s->protocol) {
+	case PROTOCOL_UDP:
+		memset(&sa->v4, 0, sizeof(sa->v4));
+		sa->s.sa_family = AF_INET;
+		sa->v4.sin_port = port;
+		memcpy(&sa->v4.sin_addr, udp_address + 1 + sizeof(uint16_t), sizeof(sa->v4.sin_addr));	// ipv4 address is 32 bits
+		return sizeof(sa->v4);
+	case PROTOCOL_UDPv6:
+		memset(&sa->v6, 0, sizeof(sa->v6));
+		sa->s.sa_family = AF_INET6;
+		sa->v6.sin6_port = port;
+		memcpy(&sa->v6.sin6_addr, udp_address + 1 + sizeof(uint16_t), sizeof(sa->v6.sin6_addr));	// ipv4 address is 128 bits
+		return sizeof(sa->v6);
+	}
+	return 0;
+}
+
+static int
+send_list_udp(struct socket_server *ss, struct socket *s, struct wb_list *list, struct socket_message *result) {
+	while (list->head) {
+		struct write_buffer * tmp = list->head;
+		union sockaddr_all sa;
+		socklen_t sasz = udp_socket_address(s, tmp->udp_address, &sa);
+		int err = sendto(s->fd, tmp->ptr, tmp->sz, 0, &sa.s, sasz);
+		if (err < 0) {
+			switch(errno) {
+			case EINTR:
+			case EAGAIN:
+				return -1;
+			}
+			fprintf(stderr, "socket-server : udp (%d) sendto error %s.\n",s->id, strerror(errno));
+			return -1;
+/*			// ignore udp sendto error
+			
+			result->opaque = s->opaque;
+			result->id = s->id;
+			result->ud = 0;
+			result->data = NULL;
+
+			return SOCKET_ERROR;
+*/
+		}
+
+		s->wb_size -= tmp->sz;
+		list->head = tmp->next;
+		write_buffer_free(ss,tmp);
+	}
+	list->tail = NULL;
+
+	return -1;
+}
+
+static int
+send_list(struct socket_server *ss, struct socket *s, struct wb_list *list, struct socket_message *result) {
+	if (s->protocol == PROTOCOL_TCP) {
+		return send_list_tcp(ss, s, list, result);
+	} else {
+		return send_list_udp(ss, s, list, result);
+	}
+}
+
 static inline int
 list_uncomplete(struct wb_list *s) {
 	struct write_buffer *wb = s->head;
@@ -500,9 +623,9 @@ send_buffer(struct socket_server *ss, struct socket *s, struct socket_message *r
 	return -1;
 }
 
-static int
-append_sendbuffer_(struct socket_server *ss, struct wb_list *s, struct request_send * request, int n) {
-	struct write_buffer * buf = MALLOC(sizeof(*buf));
+static struct write_buffer *
+append_sendbuffer_(struct socket_server *ss, struct wb_list *s, struct request_send * request, int size, int n) {
+	struct write_buffer * buf = MALLOC(size);
 	struct send_object so;
 	buf->userobject = send_object_init(ss, &so, request->buffer, request->sz);
 	buf->ptr = so.buffer+n;
@@ -517,17 +640,27 @@ append_sendbuffer_(struct socket_server *ss, struct wb_list *s, struct request_s
 		s->tail->next = buf;
 		s->tail = buf;
 	}
-	return buf->sz;
+	return buf;
+}
+
+static inline void
+append_sendbuffer_udp(struct socket_server *ss, struct socket *s, int priority, struct request_send * request, const uint8_t udp_address[UDP_ADDRESS_SIZE]) {
+	struct wb_list *wl = (priority == PRIORITY_HIGH) ? &s->high : &s->low;
+	struct write_buffer *buf = append_sendbuffer_(ss, wl, request, SIZEOF_UDPBUFFER, 0);
+	memcpy(buf->udp_address, udp_address, UDP_ADDRESS_SIZE);
+	s->wb_size += buf->sz;
 }
 
 static inline void
 append_sendbuffer(struct socket_server *ss, struct socket *s, struct request_send * request, int n) {
-	s->wb_size += append_sendbuffer_(ss, &s->high, request, n);
+	struct write_buffer *buf = append_sendbuffer_(ss, &s->high, request, SIZEOF_TCPBUFFER, n);
+	s->wb_size += buf->sz;
 }
 
 static inline void
 append_sendbuffer_low(struct socket_server *ss,struct socket *s, struct request_send * request) {
-	s->wb_size += append_sendbuffer_(ss, &s->low, request, 0);
+	struct write_buffer *buf = append_sendbuffer_(ss, &s->low, request, SIZEOF_TCPBUFFER, 0);
+	s->wb_size += buf->sz;
 }
 
 static inline int
@@ -543,7 +676,7 @@ send_buffer_empty(struct socket *s) {
 	Else append package to high (PRIORITY_HIGH) or low (PRIORITY_LOW) list.
  */
 static int
-send_socket(struct socket_server *ss, struct request_send * request, struct socket_message *result, int priority) {
+send_socket(struct socket_server *ss, struct request_send * request, struct socket_message *result, int priority, const uint8_t *udp_address) {
 	int id = request->id;
 	struct socket * s = &ss->slot[HASH_ID(id)];
 	struct send_object so;
@@ -556,30 +689,52 @@ send_socket(struct socket_server *ss, struct request_send * request, struct sock
 	}
 	assert(s->type != SOCKET_TYPE_PLISTEN && s->type != SOCKET_TYPE_LISTEN);
 	if (send_buffer_empty(s) && s->type == SOCKET_TYPE_CONNECTED) {
-		int n = write(s->fd, so.buffer, so.sz);
-		if (n<0) {
-			switch(errno) {
-			case EINTR:
-			case EAGAIN:
-				n = 0;
-				break;
-			default:
-				fprintf(stderr, "socket-server: write to %d (fd=%d) error :%s.\n",id,s->fd,strerror(errno));
-				force_close(ss,s,result);
-				return SOCKET_CLOSE;
+		if (s->protocol == PROTOCOL_TCP) {
+			int n = write(s->fd, so.buffer, so.sz);
+			if (n<0) {
+				switch(errno) {
+				case EINTR:
+				case EAGAIN:
+					n = 0;
+					break;
+				default:
+					fprintf(stderr, "socket-server: write to %d (fd=%d) error :%s.\n",id,s->fd,strerror(errno));
+					force_close(ss,s,result);
+					return SOCKET_CLOSE;
+				}
+			}
+			if (n == so.sz) {
+				so.free_func(request->buffer);
+				return -1;
+			}
+			append_sendbuffer(ss, s, request, n);	// add to high priority list, even priority == PRIORITY_LOW
+		} else {
+			// udp
+			if (udp_address == NULL) {
+				udp_address = s->p.udp_address;
+			}
+			union sockaddr_all sa;
+			socklen_t sasz = udp_socket_address(s, udp_address, &sa);
+			int n = sendto(s->fd, so.buffer, so.sz, 0, &sa.s, sasz);
+			if (n != so.sz) {
+				append_sendbuffer_udp(ss,s,priority,request,udp_address);
+			} else {
+				so.free_func(request->buffer);
 			}
 		}
-		if (n == so.sz) {
-			so.free_func(request->buffer);
-			return -1;
-		}
-		append_sendbuffer(ss, s, request, n);	// add to high priority list, even priority == PRIORITY_LOW
 		sp_write(ss->event_fd, s->fd, s, true);
 	} else {
-		if (priority == PRIORITY_LOW) {
-			append_sendbuffer_low(ss, s, request);
+		if (s->protocol == PROTOCOL_TCP) {
+			if (priority == PRIORITY_LOW) {
+				append_sendbuffer_low(ss, s, request);
+			} else {
+				append_sendbuffer(ss, s, request, 0);
+			}
 		} else {
-			append_sendbuffer(ss, s, request, 0);
+			if (udp_address == NULL) {
+				udp_address = s->p.udp_address;
+			}
+			append_sendbuffer_udp(ss,s,priority,request,udp_address);
 		}
 	}
 	return -1;
@@ -589,7 +744,7 @@ static int
 listen_socket(struct socket_server *ss, struct request_listen * request, struct socket_message *result) {
 	int id = request->id;
 	int listen_fd = request->fd;
-	struct socket *s = new_fd(ss, id, listen_fd, request->opaque, false);
+	struct socket *s = new_fd(ss, id, listen_fd, PROTOCOL_TCP, request->opaque, false);
 	if (s == NULL) {
 		goto _failed;
 	}
@@ -639,7 +794,7 @@ bind_socket(struct socket_server *ss, struct request_bind *request, struct socke
 	result->id = id;
 	result->opaque = request->opaque;
 	result->ud = 0;
-	struct socket *s = new_fd(ss, id, request->fd, request->opaque, true);
+	struct socket *s = new_fd(ss, id, request->fd, PROTOCOL_TCP, request->opaque, true);
 	if (s == NULL) {
 		result->data = NULL;
 		return SOCKET_ERROR;
@@ -719,6 +874,49 @@ has_cmd(struct socket_server *ss) {
 	return 0;
 }
 
+static void
+add_udp_socket(struct socket_server *ss, struct request_udp *udp) {
+	int id = udp->id;
+	int protocol;
+	if (udp->family == AF_INET6) {
+		protocol = PROTOCOL_UDPv6;
+	} else {
+		protocol = PROTOCOL_UDP;
+	}
+	struct socket *ns = new_fd(ss, id, udp->fd, protocol, udp->opaque, true);
+	if (ns == NULL) {
+		close(udp->fd);
+		ss->slot[HASH_ID(id)].type = SOCKET_TYPE_INVALID;
+	}
+	ns->type = SOCKET_TYPE_CONNECTED;
+	memset(ns->p.udp_address, 0, sizeof(ns->p.udp_address));
+}
+
+static int
+set_udp_address(struct socket_server *ss, struct request_setudp *request, struct socket_message *result) {
+	int id = request->id;
+	struct socket *s = &ss->slot[HASH_ID(id)];
+	if (s->type == SOCKET_TYPE_INVALID || s->id !=id) {
+		return -1;
+	}
+	int type = request->address[0];
+	if (type != s->protocol) {
+		// protocol mismatch
+		result->opaque = s->opaque;
+		result->id = s->id;
+		result->ud = 0;
+		result->data = NULL;
+
+		return SOCKET_ERROR;
+	}
+	if (type == PROTOCOL_UDP) {
+		memcpy(s->p.udp_address, request->address, 1+2+4);	// 1 type, 2 port, 4 ipv4
+	} else {
+		memcpy(s->p.udp_address, request->address, 1+2+16);	// 1 type, 2 port, 16 ipv6
+	}
+	return -1;
+}
+
 // return type
 static int
 ctrl_cmd(struct socket_server *ss, struct socket_message *result) {
@@ -749,12 +947,21 @@ ctrl_cmd(struct socket_server *ss, struct socket_message *result) {
 		result->data = NULL;
 		return SOCKET_EXIT;
 	case 'D':
-		return send_socket(ss, (struct request_send *)buffer, result, PRIORITY_HIGH);
+		return send_socket(ss, (struct request_send *)buffer, result, PRIORITY_HIGH, NULL);
 	case 'P':
-		return send_socket(ss, (struct request_send *)buffer, result, PRIORITY_LOW);
+		return send_socket(ss, (struct request_send *)buffer, result, PRIORITY_LOW, NULL);
+	case 'A': {
+		struct request_send_udp * rsu = (struct request_send_udp *)buffer;
+		return send_socket(ss, &rsu->send, result, PRIORITY_HIGH, rsu->address);
+	}
+	case 'C':
+		return set_udp_address(ss, (struct request_setudp *)buffer, result);
 	case 'T':
 		setopt_socket(ss, (struct request_setopt *)buffer);
 		return -1;
+	case 'U':
+		add_udp_socket(ss, (struct request_udp *)buffer);
+		return -1;
 	default:
 		fprintf(stderr, "socket-server: Unknown ctrl %c.\n",type);
 		return -1;
@@ -765,8 +972,8 @@ ctrl_cmd(struct socket_server *ss, struct socket_message *result) {
 
 // return -1 (ignore) when error
 static int
-forward_message(struct socket_server *ss, struct socket *s, struct socket_message * result) {
-	int sz = s->size;
+forward_message_tcp(struct socket_server *ss, struct socket *s, struct socket_message * result) {
+	int sz = s->p.size;
 	char * buffer = MALLOC(sz);
 	int n = (int)read(s->fd, buffer, sz);
 	if (n<0) {
@@ -797,9 +1004,9 @@ forward_message(struct socket_server *ss, struct socket *s, struct socket_messag
 	}
 
 	if (n == sz) {
-		s->size *= 2;
+		s->p.size *= 2;
 	} else if (sz > MIN_READ_BUFFER && n*2 < sz) {
-		s->size /= 2;
+		s->p.size /= 2;
 	}
 
 	result->opaque = s->opaque;
@@ -809,6 +1016,63 @@ forward_message(struct socket_server *ss, struct socket *s, struct socket_messag
 	return SOCKET_DATA;
 }
 
+static int
+gen_udp_address(int protocol, union sockaddr_all *sa, uint8_t * udp_address) {
+	int addrsz = 1;
+	udp_address[0] = (uint8_t)protocol;
+	if (protocol == PROTOCOL_UDP) {
+		memcpy(udp_address+addrsz, &sa->v4.sin_port, sizeof(sa->v4.sin_port));
+		addrsz += sizeof(sa->v4.sin_port);
+		memcpy(udp_address+addrsz, &sa->v4.sin_addr, sizeof(sa->v4.sin_addr));
+		addrsz += sizeof(sa->v4.sin_addr);
+	} else {
+		memcpy(udp_address+addrsz, &sa->v6.sin6_port, sizeof(sa->v6.sin6_port));
+		addrsz += sizeof(sa->v6.sin6_port);
+		memcpy(udp_address+addrsz, &sa->v6.sin6_addr, sizeof(sa->v6.sin6_addr));
+		addrsz += sizeof(sa->v6.sin6_addr);
+	}
+	return addrsz;
+}
+
+static int
+forward_message_udp(struct socket_server *ss, struct socket *s, struct socket_message * result) {
+	union sockaddr_all sa;
+	socklen_t slen = sizeof(sa);
+	int n = recvfrom(s->fd, ss->udpbuffer,MAX_UDP_PACKAGE,0,&sa.s,&slen);
+	if (n<0) {
+		switch(errno) {
+		case EINTR:
+		case EAGAIN:
+			break;
+		default:
+			// close when error
+			force_close(ss, s, result);
+			return SOCKET_ERROR;
+		}
+		return -1;
+	}
+	uint8_t * data;
+	if (slen == sizeof(sa.v4)) {
+		if (s->protocol != PROTOCOL_UDP)
+			return -1;
+		data = MALLOC(n + 1 + 2 + 4);
+		gen_udp_address(PROTOCOL_UDP, &sa, data + n);
+	} else {
+		if (s->protocol != PROTOCOL_UDPv6)
+			return -1;
+		data = MALLOC(n + 1 + 2 + 16);
+		gen_udp_address(PROTOCOL_UDPv6, &sa, data + n);
+	}
+	memcpy(data, ss->udpbuffer, n);
+
+	result->opaque = s->opaque;
+	result->id = s->id;
+	result->ud = n;
+	result->data = (char *)data;
+
+	return SOCKET_UDP;
+}
+
 static int
 report_connect(struct socket_server *ss, struct socket *s, struct socket_message *result) {
 	int error;
@@ -855,7 +1119,7 @@ report_accept(struct socket_server *ss, struct socket *s, struct socket_message
 	}
 	socket_keepalive(client_fd);
 	sp_nonblocking(client_fd);
-	struct socket *ns = new_fd(ss, id, client_fd, s->opaque, false);
+	struct socket *ns = new_fd(ss, id, client_fd, PROTOCOL_TCP, s->opaque, false);
 	if (ns == NULL) {
 		close(client_fd);
 		return 0;
@@ -941,7 +1205,17 @@ socket_server_poll(struct socket_server *ss, struct socket_message * result, int
 			break;
 		default:
 			if (e->read) {
-				int type = forward_message(ss, s, result);
+				int type;
+				if (s->protocol == PROTOCOL_TCP) {
+					type = forward_message_tcp(ss, s, result);
+				} else {
+					type = forward_message_udp(ss, s, result);
+					if (type == SOCKET_UDP) {
+						// try read again
+						--ss->event_index;
+						return SOCKET_UDP;
+					}
+				}
 				if (e->write) {
 					// Try to dispatch write message next step if write flag set.
 					e->read = false;
@@ -986,9 +1260,11 @@ open_request(struct socket_server *ss, struct request_package *req, uintptr_t op
 	int len = strlen(addr);
 	if (len + sizeof(req->u.open) > 256) {
 		fprintf(stderr, "socket-server : Invalid addr %s.\n",addr);
-		return 0;
+		return -1;
 	}
 	int id = reserve_id(ss);
+	if (id < 0)
+		return -1;
 	req->u.open.opaque = opaque;
 	req->u.open.id = id;
 	req->u.open.port = port;
@@ -1002,6 +1278,8 @@ int
 socket_server_connect(struct socket_server *ss, uintptr_t opaque, const char * addr, int port) {
 	struct request_package request;
 	int len = open_request(ss, &request, opaque, addr, port);
+	if (len < 0)
+		return -1;
 	send_request(ss, &request, 'O', sizeof(request.u.open) + len);
 	return request.u.open.id;
 }
@@ -1052,38 +1330,67 @@ socket_server_close(struct socket_server *ss, uintptr_t opaque, int id) {
 	send_request(ss, &request, 'K', sizeof(request.u.close));
 }
 
+// return -1 means failed
+// or return AF_INET or AF_INET6
+static int
+do_bind(const char *host, int port, int protocol, int *family) {
+	int fd;
+	int status;
+	int reuse = 1;
+	struct addrinfo ai_hints;
+	struct addrinfo *ai_list = NULL;
+	char portstr[16];
+	if (host == NULL || host[0] == 0) {
+		host = "0.0.0.0";	// INADDR_ANY
+	}
+	sprintf(portstr, "%d", port);
+	memset( &ai_hints, 0, sizeof( ai_hints ) );
+	ai_hints.ai_family = AF_UNSPEC;
+	if (protocol == IPPROTO_TCP) {
+		ai_hints.ai_socktype = SOCK_STREAM;
+	} else {
+		assert(protocol == IPPROTO_UDP);
+		ai_hints.ai_socktype = SOCK_DGRAM;
+	}
+	ai_hints.ai_protocol = protocol;
+
+	status = getaddrinfo( host, portstr, &ai_hints, &ai_list );
+	if ( status != 0 ) {
+		return -1;
+	}
+	*family = ai_list->ai_family;
+	fd = socket(*family, ai_list->ai_socktype, 0);
+	if (fd < 0) {
+		goto _failed_fd;
+	}
+	if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (void *)&reuse, sizeof(int))==-1) {
+		goto _failed;
+	}
+	status = bind(fd, (struct sockaddr *)ai_list->ai_addr, ai_list->ai_addrlen);
+	if (status != 0)
+		goto _failed;
+
+	freeaddrinfo( ai_list );
+	return fd;
+_failed:
+	close(fd);
+_failed_fd:
+	freeaddrinfo( ai_list );
+	return -1;
+}
+
 static int
 do_listen(const char * host, int port, int backlog) {
-	// only support ipv4
-	// todo: support ipv6 by getaddrinfo
-	uint32_t addr = INADDR_ANY;
-	if (host[0]) {
-		addr=inet_addr(host);
-	}
-	int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
+	int family = 0;
+	int listen_fd = do_bind(host, port, IPPROTO_TCP, &family);
 	if (listen_fd < 0) {
 		return -1;
 	}
-	int reuse = 1;
-	if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, (void *)&reuse, sizeof(int))==-1) {
-		goto _failed;
-	}
-
-	struct sockaddr_in my_addr;
-	memset(&my_addr, 0, sizeof(struct sockaddr_in));
-	my_addr.sin_family = AF_INET;
-	my_addr.sin_port = htons(port);
-	my_addr.sin_addr.s_addr = addr;
-	if (bind(listen_fd, (struct sockaddr *)&my_addr, sizeof(struct sockaddr)) == -1) {
-		goto _failed;
-	}
 	if (listen(listen_fd, backlog) == -1) {
-		goto _failed;
+		close(listen_fd);
+		return -1;
 	}
 	return listen_fd;
-_failed:
-	close(listen_fd);
-	return -1;
 }
 
 int 
@@ -1094,6 +1401,10 @@ socket_server_listen(struct socket_server *ss, uintptr_t opaque, const char * ad
 	}
 	struct request_package request;
 	int id = reserve_id(ss);
+	if (id < 0) {
+		close(fd);
+		return id;
+	}
 	request.u.listen.opaque = opaque;
 	request.u.listen.id = id;
 	request.u.listen.fd = fd;
@@ -1105,6 +1416,8 @@ int
 socket_server_bind(struct socket_server *ss, uintptr_t opaque, int fd) {
 	struct request_package request;
 	int id = reserve_id(ss);
+	if (id < 0)
+		return -1;
 	request.u.bind.opaque = opaque;
 	request.u.bind.id = id;
 	request.u.bind.fd = fd;
@@ -1134,3 +1447,124 @@ socket_server_userobject(struct socket_server *ss, struct socket_object_interfac
 	ss->soi = *soi;
 }
 
+// UDP
+
+int 
+socket_server_udp(struct socket_server *ss, uintptr_t opaque, const char * addr, int port) {
+	int fd;
+	int family;
+	if (port != 0 || addr != NULL) {
+		// bind
+		fd = do_bind(addr, port, IPPROTO_UDP, &family);
+		if (fd < 0) {
+			return -1;
+		}
+	} else {
+		family = AF_INET;
+		fd = socket(family, SOCK_DGRAM, 0);
+		if (fd < 0) {
+			return -1;
+		}
+	}
+	sp_nonblocking(fd);
+
+	int id = reserve_id(ss);
+	if (id < 0) {
+		close(fd);
+		return -1;
+	}
+	struct request_package request;
+	request.u.udp.id = id;
+	request.u.udp.fd = fd;
+	request.u.udp.opaque = opaque;
+	request.u.udp.family = family;
+
+	send_request(ss, &request, 'U', sizeof(request.u.udp));	
+	return id;
+}
+
+int64_t 
+socket_server_udp_send(struct socket_server *ss, int id, const struct socket_udp_address *addr, const void *buffer, int sz) {
+	struct socket * s = &ss->slot[HASH_ID(id)];
+	if (s->id != id || s->type == SOCKET_TYPE_INVALID) {
+		return -1;
+	}
+
+	struct request_package request;
+	request.u.send_udp.send.id = id;
+	request.u.send_udp.send.sz = sz;
+	request.u.send_udp.send.buffer = (char *)buffer;
+
+	const uint8_t *udp_address = (const uint8_t *)addr;
+	int addrsz;
+	switch (udp_address[0]) {
+	case PROTOCOL_UDP:
+		addrsz = 1+2+4;		// 1 type, 2 port, 4 ipv4
+		break;
+	case PROTOCOL_UDPv6:
+		addrsz = 1+2+16;	// 1 type, 2 port, 16 ipv6
+		break;
+	default:
+		return -1;
+	}
+
+	memcpy(request.u.send_udp.address, udp_address, addrsz);	
+
+	send_request(ss, &request, 'A', sizeof(request.u.send_udp.send)+addrsz);
+	return s->wb_size;
+}
+
+int
+socket_server_udp_connect(struct socket_server *ss, int id, const char * addr, int port) {
+	int status;
+	struct addrinfo ai_hints;
+	struct addrinfo *ai_list = NULL;
+	char portstr[16];
+	sprintf(portstr, "%d", port);
+	memset( &ai_hints, 0, sizeof( ai_hints ) );
+	ai_hints.ai_family = AF_UNSPEC;
+	ai_hints.ai_socktype = SOCK_DGRAM;
+	ai_hints.ai_protocol = IPPROTO_UDP;
+
+	status = getaddrinfo(addr, portstr, &ai_hints, &ai_list );
+	if ( status != 0 ) {
+		return -1;
+	}
+	struct request_package request;
+	request.u.set_udp.id = id;
+	int protocol;
+
+	if (ai_list->ai_family == AF_INET) {
+		protocol = PROTOCOL_UDP;
+	} else if (ai_list->ai_family == AF_INET6) {
+		protocol = PROTOCOL_UDPv6;
+	} else {
+		freeaddrinfo( ai_list );
+		return -1;
+	}
+
+	int addrsz = gen_udp_address(protocol, (union sockaddr_all *)ai_list->ai_addr, request.u.set_udp.address);
+
+	freeaddrinfo( ai_list );
+
+	send_request(ss, &request, 'C', sizeof(request.u.set_udp) - sizeof(request.u.set_udp.address) +addrsz);
+
+	return 0;
+}
+
+const struct socket_udp_address *
+socket_server_udp_address(struct socket_server *ss, struct socket_message *msg, int *addrsz) {
+	uint8_t * address = (uint8_t *)(msg->data + msg->ud);
+	int type = address[0];
+	switch(type) {
+	case PROTOCOL_UDP:
+		*addrsz = 1+2+4;
+		break;
+	case PROTOCOL_UDPv6:
+		*addrsz = 1+2+16;
+		break;
+	default:
+		return NULL;
+	}
+	return (const struct socket_udp_address *)address;
+}

skynet-src/socket_server.h

@@ -9,6 +9,7 @@
 #define SOCKET_ACCEPT 3
 #define SOCKET_ERROR 4
 #define SOCKET_EXIT 5
+#define SOCKET_UDP 6
 
 struct socket_server;
 
@@ -36,8 +37,22 @@ int socket_server_listen(struct socket_server *, uintptr_t opaque, const char *
 int socket_server_connect(struct socket_server *, uintptr_t opaque, const char * addr, int port);
 int socket_server_bind(struct socket_server *, uintptr_t opaque, int fd);
 
+// for tcp
 void socket_server_nodelay(struct socket_server *, int id);
 
+struct socket_udp_address;
+
+// create an udp socket handle, attach opaque with it . udp socket don't need call socket_server_start to recv message
+// if port != 0, bind the socket . if addr == NULL, bind ipv4 0.0.0.0 . If you want to use ipv6, addr can be "::" and port 0.
+int socket_server_udp(struct socket_server *, uintptr_t opaque, const char * addr, int port);
+// set default dest address, return 0 when success
+int socket_server_udp_connect(struct socket_server *, int id, const char * addr, int port);
+// If the socket_udp_address is NULL, use last call socket_server_udp_connect address instead
+// You can also use socket_server_send 
+int64_t socket_server_udp_send(struct socket_server *, int id, const struct socket_udp_address *, const void *buffer, int sz);
+// extract the address of the message, struct socket_message * should be SOCKET_UDP
+const struct socket_udp_address * socket_server_udp_address(struct socket_server *, struct socket_message *, int *addrsz);
+
 struct socket_object_interface {
 	void * (*buffer)(void *);
 	int (*size)(void *);

test/testudp.lua

@@ -0,0 +1,26 @@
+local skynet = require "skynet"
+local socket = require "socket"
+
+local function server()
+	local host
+	host = socket.udp(function(data, sz, from)
+		local str = skynet.tostring(data,sz)	-- skynet.tostring should call only once, because it will free the pointer data
+		print("server recv", str, socket.udp_address(from))
+		socket.sendto(host, from, "OK " .. str)
+	end , "127.0.0.1", 8765)	-- bind an address
+end
+
+local function client()
+	local c = socket.udp(function(data, sz, from)
+		print("client recv", skynet.tostring(data,sz), socket.udp_address(from))
+	end)
+	socket.udp_connect(c, "127.0.0.1", 8765)
+	for i=1,20 do
+		socket.write(c, "hello " .. i)	-- write to the address by udp_connect binding
+	end
+end
+
+skynet.start(function()
+	skynet.fork(server)
+	skynet.fork(client)
+end)