http_reader: improve robustness of header reading
[ruby-mogilefs-client.git] / lib / mogilefs / http_reader.rb
blob8a0337ee2fb6c1a956f1f8315695ef97fbf11ecf
1 # -*- encoding: binary -*-
2 # internal implementation details here, do not rely on them in your code
4 # This class is needed because Net::HTTP streaming is still inefficient
5 # for reading huge response bodies over fast LANs.
6 class MogileFS::HTTPReader < MogileFS::Socket
7   attr_accessor :content_length, :uri
9   # backwards compat, if anybody cares
10   alias mogilefs_size content_length # :nodoc:
12   # this may OOM your system on large files
13   def to_s
14     buf = ""
15     read(@content_length, buf)
16     return buf if buf.size == @content_length
18     raise MogileFS::SizeMismatchError,
19           "read=#{buf.size} bytes, expected=#@content_length from #@uri", []
20   end
22   def stream_to(dest)
23     rv = MogileFS.io.copy_stream(self, dest)
24     return rv if rv == @content_length
25     raise MogileFS::SizeMismatchError,
26           "read=#{rv} bytes, expected=#@content_length from #@uri", []
27   end
29   def self.first(paths, timeout, count = nil, offset = nil)
30     errors = nil
31     if offset || count
32       offset ||= 0
33       range_end = count ? offset + count - 1 : ""
34       range = "Range: bytes=#{offset}-#{range_end}\r\n"
35     end
37     paths.each do |path|
38       begin
39         sock = try(path, timeout, range) and return sock
40       rescue => e
41         errors ||= []
42         errors << "#{path} - #{e.message} (#{e.class})"
43       end
44     end
45     raise MogileFS::Error,
46           "all paths failed with GET: #{errors.join(', ')}", []
47   end
49   # given a path, this returns a readable socket with ready data from the
50   # body of the response.
51   def self.try(path, timeout, range) # :nodoc:
52     uri = URI.parse(path)
53     expire_at = Time.now + timeout
54     sock = tcp(uri.host, uri.port, timeout)
55     buf = "GET #{uri.request_uri} HTTP/1.0\r\n#{range}\r\n" # no chunking
56     sock.timed_write(buf, timeout)
58     begin
59       raise MogileFS::Timeout if Time.now > expire_at
60       sock.timed_peek(2048, buf, timeout) or
61         raise MogileFS::InvalidResponseError, "EOF while reading header", []
62     end until /\r\n\r\n/ =~ buf
64     head, _ = buf.split(/\r\n\r\n/, 2)
65     if ((range && head =~ %r{\AHTTP/\d+\.\d+\s+206\s*}) ||
66         (!range && head =~ %r{\AHTTP/\d+\.\d+\s+200\s*})) &&
67        head =~ %r{^Content-Length:\s*(\d+)}i
68       sock.content_length = $1.to_i
69       sock.uri = uri
70       sock.timed_read(head.bytesize + 4, buf, 0)
71       return sock
72     end
73     msg = range ? "Expected 206 w/#{range.strip}: " : "header="
74     msg << head.inspect
75     raise MogileFS::InvalidResponseError, msg, []
76   rescue
77     sock.close if sock
78     raise
79   end
80 end