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