lua/lib/socket.lua

   1 -- transparent non-blocking writes with buffers!
   2
   3 local socket_write_raw = socket_write
   4 local socket_unix_connect_raw = socket_unix_connect
   5 local socket_unix_bind_raw = socket_unix_bind
   6 local socket_tcp_connect_raw = socket_tcp_connect
   7 local socket_tcp_bind_raw = socket_tcp_bind
   8 local socket_udp_connect_raw = socket_udp_connect
   9 local socket_udp_bind_raw = socket_udp_bind
  10 local socket_close_raw = socket_close
  11
  12 local sockets = {}
  13
  14 local function socket_handler(socket, event, tag, ...)
  15   if event == "flush" then
  16     local buf = sockets[socket].writebuf
  17     local ret = socket_write_raw(socket, buf)
  18
  19     if ret == -1 then
  20       return -- close will be called
  21     end
  22
  23     sockets[socket].writebuf = buf:sub(ret + 1)
  24
  25     if sockets[socket].closing and sockets[socket].writebuf == "" then
  26       socket_close(socket)
  27     end
  28
  29     -- no reason to tell the caller
  30     return
  31   elseif event == "accept" then
  32     local newsocket = ...
  33     socket_new(newsocket, sockets[socket].handler)
  34   end
  35
  36   sockets[socket].handler(socket, event, tag, ...)
  37
  38   if event == "close" then
  39     sockets[socket] = nil
  40   end
  41 end
  42
  43 function socket_new(socket, handler)
  44   sockets[socket] = { writebuf = "", handler = handler }
  45 end
  46
  47 function socket_unix_connect(path, handler, tag)
  48   local connected, socket = socket_unix_connect_raw(path, socket_handler, tag)
  49   if connected == nil then
  50     return nil
  51   end
  52
  53   socket_new(socket, handler)
  54   if connected then
  55     socket_handler(socket, "connect", tag)
  56   end
  57
  58   return socket
  59 end
  60
  61 function socket_unix_bind(path, handler, tag)
  62   local socket = socket_unix_bind_raw(path, socket_handler, tag)
  63   if not socket then
  64     return nil
  65   end
  66
  67   socket_new(socket, handler)
  68   return socket
  69 end
  70
  71 local function socket_ip_connect(fn, address, port, handler, tag)
  72   local connected, socket = fn(address, port, socket_handler, tag)
  73
  74   if connected == nil then
  75     return nil
  76   end
  77
  78   socket_new(socket, handler)
  79   if connected then
  80     socket_handler(socket, "connect", tag)
  81   end
  82
  83   return socket
  84 end
  85
  86 local function socket_ip_bind(fn, address, port, handler, tag)
  87   local socket = fn(address, port, socket_handler, tag)
  88   if not socket then
  89     return nil
  90   end
  91
  92   socket_new(socket, handler)
  93   return socket
  94 end
  95
  96 function socket_tcp_bind(address, port, handler, tag)
  97   return socket_ip_bind(socket_tcp_bind_raw, address, port, handler, tag)
  98 end
  99
 100 function socket_tcp_connect(address, port, handler, tag)
 101   return socket_ip_connect(socket_tcp_connect_raw, address, port, handler, tag)
 102 end
 103
 104 function socket_udp_bind(address, port, handler, tag)
 105   return socket_ip_bind(socket_udp_bind_raw, address, port, handler, tag)
 106 end
 107
 108 function socket_udp_connect(address, port, handler, tag)
 109   return socket_ip_connect(socket_udp_connect_raw, address, port, handler, tag)
 110 end
 111
 112 function socket_write(socket, data)
 113   if sockets[socket].writebuf == "" then
 114     local ret = socket_write_raw(socket, data)
 115     if ret == -1 then
 116       return false -- close will be called regardless
 117     end
 118
 119     if ret == data:len() then
 120       return true
 121     end
 122
 123     -- lua strings start at 1
 124     sockets[socket].writebuf = data:sub(ret + 1)
 125   else
 126     sockets[socket].writebuf = sockets[socket].writebuf .. data
 127   end
 128
 129   return true
 130 end
 131
 132 function socket_close(socket, flush)
 133   if whenbufferempty and sockets[socket].writebuf ~= "" then
 134     sockets[socket].closing = true
 135     return
 136   end
 137
 138   return socket_close_raw(socket)
 139 end