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         if iter then
  32                 for obj in iter do
  33                         tbl[#tbl+1] = obj
  34                 end
  35         end
  36         return tbl
  37 end
  38
  39 local meta = {}
  40
  41 function meta.is_socket(self)
  42         return (getmetatable(self) == socket)
  43 end
  44
  45 function meta.is_tls_socket(self)
  46         return (getmetatable(self) == tls_socket)
  47 end
  48
  49 function meta.is_file(self)
  50         return (getmetatable(self) == file)
  51 end
  52
  53 function meta.readall(self, len)
  54         local block, code, msg = self:read(len or BUFFERSIZE)
  55
  56         if not block then
  57                 return nil, code, msg, ""
  58         elseif #block == 0 then
  59                 return "", nil, nil, ""
  60         end
  61
  62         local data, total = {block}, #block
  63
  64         while not len or len > total do
  65                 block, code, msg = self:read(len and (len - total) or BUFFERSIZE)
  66
  67                 if not block then
  68                         return nil, code, msg, table.concat(data)
  69                 elseif #block == 0 then
  70                         break
  71                 end
  72
  73                 data[#data+1], total = block, total + #block
  74         end
  75
  76         local data = #data > 1 and table.concat(data) or data[1]
  77         return data, nil, nil, data
  78 end
  79 meta.recvall = meta.readall
  80
  81 function meta.writeall(self, data)
  82         data = tostring(data)
  83         local sent, code, msg = self:write(data)
  84
  85         if not sent then
  86                 return nil, code, msg, 0
  87         end
  88
  89         local total = sent
  90
  91         while total < #data do
  92                 sent, code, msg = self:write(data, total)
  93
  94                 if not sent then
  95                         return nil, code, msg, total
  96                 end
  97
  98                 total = total + sent
  99         end
 100
 101         return total, nil, nil, total
 102 end
 103 meta.sendall = meta.writeall
 104
 105 function meta.linesource(self, limit)
 106         limit = limit or BUFFERSIZE
 107         local buffer = ""
 108         local bpos = 0
 109         return function(flush)
 110                 local line, endp, _
 111
 112                 if flush then
 113                         line = buffer:sub(bpos + 1)
 114                         buffer = type(flush) == "string" and flush or ""
 115                         bpos = 0
 116                         return line
 117                 end
 118
 119                 while not line do
 120                         _, endp, line = buffer:find("(.-)\r?\n", bpos + 1)
 121                         if line then
 122                                 bpos = endp
 123                                 return line
 124                         elseif #buffer < limit + bpos then
 125                                 local newblock, code, msg = self:read(limit + bpos - #buffer)
 126                                 if not newblock then
 127                                         return nil, code, msg
 128                                 elseif #newblock == 0 then
 129                                         return nil
 130                                 end
 131                                 buffer = buffer:sub(bpos + 1) .. newblock
 132                                 bpos = 0
 133                         else
 134                                 return nil, 0
 135                         end
 136                 end
 137         end
 138 end
 139
 140 function meta.blocksource(self, bs, limit)
 141         bs = bs or BUFFERSIZE
 142         return function()
 143                 local toread = bs
 144                 if limit then
 145                         if limit < 1 then
 146                                 return nil
 147                         elseif limit < toread then
 148                                 toread = limit
 149                         end
 150                 end
 151
 152                 local block, code, msg = self:read(toread)
 153
 154                 if not block then
 155                         return nil, code, msg
 156                 elseif #block == 0 then
 157                         return nil
 158                 else
 159                         if limit then
 160                                 limit = limit - #block
 161                         end
 162
 163                         return block
 164                 end
 165         end
 166 end
 167
 168 function meta.sink(self, close)
 169         return function(chunk, src_err)
 170                 if not chunk and not src_err and close then
 171                         if self.shutdown then
 172                                 self:shutdown()
 173                         end
 174                         self:close()
 175                 elseif chunk and #chunk > 0 then
 176                         return self:writeall(chunk)
 177                 end
 178                 return true
 179         end
 180 end
 181
 182 function meta.copy(self, fdout, size)
 183         local source = self:blocksource(nil, size)
 184         local sink = fdout:sink()
 185         local sent, chunk, code, msg = 0
 186
 187         repeat
 188                 chunk, code, msg = source()
 189                 sink(chunk, code, msg)
 190                 sent = chunk and (sent + #chunk) or sent
 191         until not chunk
 192         return not code and sent or nil, code, msg, sent
 193 end
 194
 195 function meta.copyz(self, fd, size)
 196         local sent, lsent, code, msg = 0
 197         local splicable
 198
 199         if not ZBUG and self:is_file() then
 200                 local ftype = self:stat("type")
 201                 if nixio.sendfile and fd:is_socket() and ftype == "reg" then
 202                         repeat
 203                                 lsent, code, msg = nixio.sendfile(fd, self, size or ZIOBLKSIZE)
 204                                 if lsent then
 205                                         sent = sent + lsent
 206                                         size = size and (size - lsent)
 207                                 end
 208                         until (not lsent or lsent == 0 or (size and size == 0))
 209                         if lsent or (not lsent and sent == 0 and
 210                          code ~= nixio.const.ENOSYS and code ~= nixio.const.EINVAL) then
 211                                 return lsent and sent, code, msg, sent
 212                         end
 213                 elseif nixio.splice and not fd:is_tls_socket() and ftype == "fifo" then
 214                         splicable = true
 215                 end
 216         end
 217
 218         if nixio.splice and fd:is_file() and not splicable then
 219                 splicable = not self:is_tls_socket() and fd:stat("type") == "fifo"
 220         end
 221
 222         if splicable then
 223                 repeat
 224                         lsent, code, msg = nixio.splice(self, fd, size or ZIOBLKSIZE)
 225                         if lsent then
 226                                 sent = sent + lsent
 227                                 size = size and (size - lsent)
 228                         end
 229                 until (not lsent or lsent == 0 or (size and size == 0))
 230                 if lsent or (not lsent and sent == 0 and
 231                  code ~= nixio.const.ENOSYS and code ~= nixio.const.EINVAL) then
 232                         return lsent and sent, code, msg, sent
 233                 end
 234         end
 235
 236         return self:copy(fd, size)
 237 end
 238
 239 function tls_socket.close(self)
 240         return self.socket:close()
 241 end
 242
 243 function tls_socket.getsockname(self)
 244         return self.socket:getsockname()
 245 end
 246
 247 function tls_socket.getpeername(self)
 248         return self.socket:getpeername()
 249 end
 250
 251 function tls_socket.getsockopt(self, ...)
 252         return self.socket:getsockopt(...)
 253 end
 254 tls_socket.getopt = tls_socket.getsockopt
 255
 256 function tls_socket.setsockopt(self, ...)
 257         return self.socket:setsockopt(...)
 258 end
 259 tls_socket.setopt = tls_socket.setsockopt
 260
 261 for k, v in pairs(meta) do
 262         file[k] = v
 263         socket[k] = v
 264         tls_socket[k] = v
 265 end