lib/sunshowers/io.rb

   1 # -*- encoding: binary -*-
   2
   3 module Sunshowers
   4
   5   # Wraps IO objects and provides transparent framing of #gets and
   6   # #write methods.
   7   #
   8   # It is compatible with core Ruby IO objects, Revactor and
   9   # Rainbows::Fiber::IO objects.
  10   class IO < Struct.new(:to_io, :buf)
  11
  12     # most Web Socket messages from clients are expected to be small text,
  13     # so we only read 128 bytes off the socket at a time
  14     RD_SIZE = 128
  15
  16     # maximum size of a UTF-8 buffer we'll allow in memory (16K)
  17     # this is a soft limit and may be offset by the value of RD_SIZE
  18     MAX_UTF8_SIZE = 1024 * 16
  19
  20     # maximum size of a binary buffer we'll allow in memory (112K)
  21     # this is a soft limit and may be offset by the value of RD_SIZE
  22     MAX_BINARY_SIZE = 1024 * 112
  23
  24     # Web Sockets usually uses UTF-8 when interfacing with the client
  25     # Ruby Sockets always return strings of Encoding::Binary under 1.9
  26     ENC = defined?(Encoding::UTF_8) ? Encoding::UTF_8 : nil
  27
  28     Z = "" # :nodoc:
  29
  30     # Wraps the given +io+ with Web Sockets-compatible framing.
  31     #
  32     #   TCPSocket.new('example.com', 80)
  33     #   io = Sunshowers::IO.new(socket)
  34     def initialize(io, buffer = Z.dup)
  35       super
  36     end
  37
  38     # iterates through each message until a client closes the connection
  39     # or block the issues a break/return
  40     def each(&block)
  41       while str = gets
  42         yield str
  43       end
  44       self
  45     end
  46
  47     # Retrieves the next record, returns nil if client closes the connection.
  48     # The record may be either a UTF-8 or binary String, under
  49     # Ruby 1.9, the String encoding will be set appropriately.
  50     def gets
  51       begin
  52         unless buf.empty?
  53           buf.gsub!(/\A\x00(.*?)\xff/m, Z) and return utf8!($1)
  54           rv = read_binary and return rv
  55           buf.size > MAX_UTF8_SIZE and
  56             raise ProtocolError, "buffer too large #{buf.size}"
  57         end
  58         buf << read
  59       rescue EOFError
  60         return
  61       end while true
  62     end
  63
  64     def syswrite(buf) # :nodoc:
  65       to_io.write(buf)
  66       rescue EOFError,Errno::ECONNRESET,Errno::EPIPE,
  67              Errno::EINVAL,Errno::EBADF => e
  68         raise ClientShutdown, e.message, []
  69     end
  70
  71     # writes a UTF-8 frame containing +buf+.  We do not validate that
  72     # +buf+ contains valid UTF-8, we assume you know what you are
  73     # doing if you call this method directly.
  74     def write_utf8(buf)
  75       syswrite("\0#{binary(buf)}\xff")
  76     end
  77
  78     # writes a binary frame containing +buf+
  79     def write_binary(buf)
  80       buf = binary(buf)
  81       n = buf.size
  82       length = []
  83       begin
  84         length.unshift((n % 128) | 0x80)
  85       end while (n /= 128) > 0
  86       length[-1] ^= 0x80
  87
  88       syswrite("\x80#{length.pack("C*")}#{buf}")
  89     end
  90
  91     if ENC.nil?
  92       # :stopdoc:
  93       U_STAR = "U*"
  94       def valid_utf8?(buf)
  95         buf.unpack(U_STAR)
  96         true
  97         rescue ArgumentError
  98           false
  99       end
 100
 101       def write(buf)
 102         valid_utf8?(buf) ? write_utf8(buf) : write_binary(buf)
 103       end
 104       # :startdoc:
 105     else
 106
 107       # Writes out +buf+ as a Web Socket frame.  If +buf+ encoding is UTF-8,
 108       # then it will be framed as UTF-8, otherwise it will be framed as
 109       # binary with an explicit length set.
 110       def write(buf)
 111         buf.encoding == ENC ? write_utf8(buf) : write_binary(buf)
 112       end
 113     end
 114
 115     # :stopdoc:
 116     def read(size = nil)
 117       i = to_io
 118       # read with no args for Revactor compat
 119       i.respond_to?(:readpartial) ?
 120         i.readpartial(size.nil? ? RD_SIZE : size) :
 121         i.read(size.nil? ? nil : size)
 122     end
 123
 124     if ENC.nil?
 125       def utf8!(buf)
 126         valid_utf8?(buf) or raise ProtocolError, "not UTF-8: #{buf.inspect}"
 127         buf
 128       end
 129       def binary(buf); buf; end
 130     else
 131       def utf8!(buf)
 132         buf.force_encoding(ENC)
 133         buf.valid_encoding? or raise ProtocolError, "not UTF-8: #{buf.inspect}"
 134         buf
 135       end
 136       def binary(buf)
 137         buf.encoding == Encoding::BINARY ?
 138           buf : buf.dup.force_encoding(Encoding::BINARY)
 139       end
 140     end
 141
 142     if Z.respond_to?(:ord)
 143       def ord(byte_str); byte_str.ord; end
 144     else
 145       def ord(byte_str); byte_str; end
 146     end
 147
 148     def read_binary
 149       (ord(buf[0]) & 0x80) == 0x80 or return
 150
 151       i = 1
 152       b = length = 0
 153       begin
 154         buf << read while (b = buf[i]).nil?
 155         b = ord(b)
 156         length = length * 128 + (b & 0x7f)
 157         i += 1
 158       end while (b & 0x80) != 0
 159
 160       length > MAX_BINARY_SIZE and
 161         raise ProtocolError, "chunk too large: #{length} bytes"
 162
 163       to_read = length - buf.size + i
 164       while to_read > 0
 165         buf << (tmp = read)
 166         to_read -= tmp.size
 167       end
 168       rv = buf[i, length]
 169       buf.replace(buf[i+length, buf.size])
 170       rv
 171     end
 172
 173     # :startdoc:
 174   end
 175 end