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