util.lua 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270
  1. --[[
  2. nixio - Linux I/O library for lua
  3. Copyright 2009 Steven Barth <steven@midlink.org>
  4. Licensed under the Apache License, Version 2.0 (the "License");
  5. you may not use this file except in compliance with the License.
  6. You may obtain a copy of the License at
  7. http://www.apache.org/licenses/LICENSE-2.0
  8. $Id$
  9. ]]--
  10. local table = require "table"
  11. local nixio = require "nixio"
  12. local getmetatable, assert, pairs, type = getmetatable, assert, pairs, type
  13. local tostring = tostring
  14. module "nixio.util"
  15. local BUFFERSIZE = nixio.const.buffersize
  16. local ZIOBLKSIZE = 65536
  17. local socket = nixio.meta_socket
  18. local tls_socket = nixio.meta_tls_socket
  19. local file = nixio.meta_file
  20. local uname = nixio.uname()
  21. local ZBUG = uname.sysname == "Linux" and uname.release:sub(1, 3) == "2.4"
  22. function consume(iter, append)
  23. local tbl = append or {}
  24. if iter then
  25. for obj in iter do
  26. tbl[#tbl+1] = obj
  27. end
  28. end
  29. return tbl
  30. end
  31. local meta = {}
  32. function meta.is_socket(self)
  33. return (getmetatable(self) == socket)
  34. end
  35. function meta.is_tls_socket(self)
  36. return (getmetatable(self) == tls_socket)
  37. end
  38. function meta.is_file(self)
  39. return (getmetatable(self) == file)
  40. end
  41. function meta.readall(self, len)
  42. local block, code, msg = self:read(len or BUFFERSIZE)
  43. if not block then
  44. return nil, code, msg, ""
  45. elseif #block == 0 then
  46. return "", nil, nil, ""
  47. end
  48. local data, total = {block}, #block
  49. while not len or len > total do
  50. block, code, msg = self:read(len and (len - total) or BUFFERSIZE)
  51. if not block then
  52. return nil, code, msg, table.concat(data)
  53. elseif #block == 0 then
  54. break
  55. end
  56. data[#data+1], total = block, total + #block
  57. end
  58. local data = #data > 1 and table.concat(data) or data[1]
  59. return data, nil, nil, data
  60. end
  61. meta.recvall = meta.readall
  62. function meta.writeall(self, data)
  63. data = tostring(data)
  64. local sent, code, msg = self:write(data)
  65. if not sent then
  66. return nil, code, msg, 0
  67. end
  68. local total = sent
  69. while total < #data do
  70. sent, code, msg = self:write(data, total)
  71. if not sent then
  72. return nil, code, msg, total
  73. end
  74. total = total + sent
  75. end
  76. return total, nil, nil, total
  77. end
  78. meta.sendall = meta.writeall
  79. function meta.linesource(self, limit)
  80. limit = limit or BUFFERSIZE
  81. local buffer = ""
  82. local bpos = 0
  83. return function(flush)
  84. local line, endp, _
  85. if flush then
  86. line = buffer:sub(bpos + 1)
  87. buffer = type(flush) == "string" and flush or ""
  88. bpos = 0
  89. return line
  90. end
  91. while not line do
  92. _, endp, line = buffer:find("(.-)\r?\n", bpos + 1)
  93. if line then
  94. bpos = endp
  95. return line
  96. elseif #buffer < limit + bpos then
  97. local newblock, code, msg = self:read(limit + bpos - #buffer)
  98. if not newblock then
  99. return nil, code, msg
  100. elseif #newblock == 0 then
  101. return nil
  102. end
  103. buffer = buffer:sub(bpos + 1) .. newblock
  104. bpos = 0
  105. else
  106. return nil, 0
  107. end
  108. end
  109. end
  110. end
  111. function meta.blocksource(self, bs, limit)
  112. bs = bs or BUFFERSIZE
  113. return function()
  114. local toread = bs
  115. if limit then
  116. if limit < 1 then
  117. return nil
  118. elseif limit < toread then
  119. toread = limit
  120. end
  121. end
  122. local block, code, msg = self:read(toread)
  123. if not block then
  124. return nil, code, msg
  125. elseif #block == 0 then
  126. return nil
  127. else
  128. if limit then
  129. limit = limit - #block
  130. end
  131. return block
  132. end
  133. end
  134. end
  135. function meta.sink(self, close)
  136. return function(chunk, src_err)
  137. if not chunk and not src_err and close then
  138. if self.shutdown then
  139. self:shutdown()
  140. end
  141. self:close()
  142. elseif chunk and #chunk > 0 then
  143. return self:writeall(chunk)
  144. end
  145. return true
  146. end
  147. end
  148. function meta.copy(self, fdout, size)
  149. local source = self:blocksource(nil, size)
  150. local sink = fdout:sink()
  151. local sent, chunk, code, msg = 0
  152. repeat
  153. chunk, code, msg = source()
  154. sink(chunk, code, msg)
  155. sent = chunk and (sent + #chunk) or sent
  156. until not chunk
  157. return not code and sent or nil, code, msg, sent
  158. end
  159. function meta.copyz(self, fd, size)
  160. local sent, lsent, code, msg = 0
  161. local splicable
  162. if not ZBUG and self:is_file() then
  163. local ftype = self:stat("type")
  164. if nixio.sendfile and fd:is_socket() and ftype == "reg" then
  165. repeat
  166. lsent, code, msg = nixio.sendfile(fd, self, size or ZIOBLKSIZE)
  167. if lsent then
  168. sent = sent + lsent
  169. size = size and (size - lsent)
  170. end
  171. until (not lsent or lsent == 0 or (size and size == 0))
  172. if lsent or (not lsent and sent == 0 and
  173. code ~= nixio.const.ENOSYS and code ~= nixio.const.EINVAL) then
  174. return lsent and sent, code, msg, sent
  175. end
  176. elseif nixio.splice and not fd:is_tls_socket() and ftype == "fifo" then
  177. splicable = true
  178. end
  179. end
  180. if nixio.splice and fd:is_file() and not splicable then
  181. splicable = not self:is_tls_socket() and fd:stat("type") == "fifo"
  182. end
  183. if splicable then
  184. repeat
  185. lsent, code, msg = nixio.splice(self, fd, size or ZIOBLKSIZE)
  186. if lsent then
  187. sent = sent + lsent
  188. size = size and (size - lsent)
  189. end
  190. until (not lsent or lsent == 0 or (size and size == 0))
  191. if lsent or (not lsent and sent == 0 and
  192. code ~= nixio.const.ENOSYS and code ~= nixio.const.EINVAL) then
  193. return lsent and sent, code, msg, sent
  194. end
  195. end
  196. return self:copy(fd, size)
  197. end
  198. if tls_socket then
  199. function tls_socket.close(self)
  200. return self.socket:close()
  201. end
  202. function tls_socket.getsockname(self)
  203. return self.socket:getsockname()
  204. end
  205. function tls_socket.getpeername(self)
  206. return self.socket:getpeername()
  207. end
  208. function tls_socket.getsockopt(self, ...)
  209. return self.socket:getsockopt(...)
  210. end
  211. tls_socket.getopt = tls_socket.getsockopt
  212. function tls_socket.setsockopt(self, ...)
  213. return self.socket:setsockopt(...)
  214. end
  215. tls_socket.setopt = tls_socket.setsockopt
  216. end
  217. for k, v in pairs(meta) do
  218. file[k] = v
  219. socket[k] = v
  220. if tls_socket then
  221. tls_socket[k] = v
  222. end
  223. end