libs/nixio/lua/nixio/util.lua

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