# -*- encoding: binary -*-
require 'digest/md5'
require 'rack'

# Rack application for handling HTTP PUT/DELETE/MKCOL operations needed
# for a MogileFS storage server.  GET and HEAD requests are handled by
# Rack::File and Rack::Head respectively.
class MogstoredRack

  # Basic usage in rackup config file (config.ru):
  #
  #    require "mogstored_rack"
  #    run MogstoredRack.new(path = "/var/mogdata", options = {})
  #
  # The +options+ hash accepts several configuration options:
  #
  # [:open_flags => additional integer flags (IO::* constants)]
  #
  #   You man specify IO::SYNC or IO::DIRECT here if your system supports it.
  #   default: 0 (IO::RDWR|IO::CREAT|IO::EXCL are always enforced)
  #
  # [:fsync => true or false]
  #
  #   Runs fsync(2) on the file and parent directory after PUT
  #   default: false
  #
  # [:io_size => Integer]
  #
  #   The I/O chunk size to use for writing PUT requests
  #   some systems may benefit from larger or smaller values
  #   default: 0x100000 (1 megabyte)
  #
  # [:put_perms => Integer (octal mask)]
  #
  #   permissions for newly-created files via PUT
  #   default: (~File.umask & 0666)
  #
  # [:mkcol_perms => Integer (octal mask)]
  #
  #   permissions for newly-created directories via MKCOL
  #   default: (~File.umask & 0777)
  #
  def initialize(root = "/var/mogdata", opts = {})
    @root = File.expand_path(root)
    @io_size = opts[:io_size] || 0x100000
    @rack_file = opts[:app] || Rack::Head.new(Rack::File.new(@root))
    @fsync = !! opts[:fsync]
    @put_perms = opts[:put_perms] || (~File.umask & 0666)
    @mkcol_perms = opts[:mkcol_perms] || (~File.umask & 0777)
    @reread_verify = !! opts[:reread_verify] # unsupported
    @open_flags = opts[:open_flags] || 0
    @open_flags |= IO::RDWR | IO::CREAT
  end

  def call(env) # :nodoc:
    case env["REQUEST_METHOD"]
    when "GET", "HEAD"
      case env["PATH_INFO"]
      when "/"
        r(200, "") # "mogadm check" uses this
      else
        @rack_file.call(env)
      end
    when "PUT"
      put(env)
    when "DELETE"
      delete(env)
    when "MKCOL"
      mkcol(env)
    else
      r(405, "unsupported method", env)
    end
    rescue Errno::EPERM, Errno::EACCES => err
      r(403, "#{err.message} (#{err.class})", env)
    rescue => err
      r(500, "#{err.message} (#{err.class})", env)
  end

  def mkcol(env) # :nodoc:
    path = server_path(env) or return r(400)
    Dir.mkdir(path, @mkcol_perms)
    r(204)
    rescue Errno::EEXIST # succeed (204) on race condition
      File.directory?(path) ? r(204) : r(409)
  end

  def delete(env) # :nodoc:
    path = server_path(env) or return r(400)
    File.exist?(path) or return r(404)
    File.directory?(path) ? Dir.rmdir(path) : File.unlink(path)
    r(204)
    rescue Errno::ENOENT # return 404 on race condition
      File.exist?(path) ? r(500) : r(404)
  end

  def put_range(range, env, path) # :nodoc:
    %r{\A\s*bytes\s+(\d+)-(\d+)/\*\s*\z} =~ range or
      return r(400, "Bad range", env)
    clen = env["CONTENT_LENGTH"] or
               return r(400, "Content-Length required for Content-Range")
    off_out = $1.to_i
    len = $2.to_i - off_out + 1
    len == clen.to_i or
             return r(400,
                      "Bad range, Content-Range: #{range} does not match\n" \
                      "Content-Length: #{clen.inspect}", env)
    File.open(path, @open_flags, 0600) { |fp| put_write(env, fp, off_out) }
  end

  def put_write(env, fp, offset = nil) # :nodoc:
    fp.binmode
    fp.sync = true
    fp.seek(offset) if offset
    received_md5 = put_loop(env, fp)
    if err = content_md5_fail?(env, received_md5)
      File.unlink(fp.path) unless offset
      return err
    end
    fp.chmod(@put_perms)
    if @fsync
      fp.fsync
      File.open(File.dirname(fp.path)) { |io| io.fsync }
    end
    r(201)
  end

  def put(env) # :nodoc:
    path = server_path(env) or return r(400)
    range = env["HTTP_CONTENT_RANGE"] and
      return put_range(range, env, path)
    File.open(path, @open_flags|IO::TRUNC, 0600) { |fp| put_write(env, fp) }
    rescue Errno::ENOENT
      r(403)
  end

  def put_loop(env, dst) # :nodoc:
    src = env["rack.input"]
    buf = ""
    md5 = Digest::MD5.new
    while src.read(@io_size, buf)
      md5.update(buf)
      dst.write(buf)
    end
    [ md5.digest ].pack('m').strip!
  end

  def server_path(env) # :nodoc:
    path = env['PATH_INFO'].squeeze('/')
    path.split(%r{/}).include?("..") and return false
    "#@root#{path}"
  end

  # returns a plain-text HTTP response
  def r(code, msg = nil, env = nil) # :nodoc:
    if env && logger = env["rack.logger"]
      logger.warn("#{env['REQUEST_METHOD']} #{env['PATH_INFO']} " \
                  "#{code} #{msg.inspect}")
    end
    if Rack::Utils::STATUS_WITH_NO_ENTITY_BODY.include?(code)
      [ code, {}, [] ]
    else
      msg ||= Rack::Utils::HTTP_STATUS_CODES[code] || ""
      msg += "\n" if msg.size > 0
      [ code,
        { 'Content-Type' => 'text/plain', 'Content-Length' => msg.size.to_s },
        [ msg ] ]
    end
  end

  # Tries to detect network corruption by verifying the client-supplied
  # Content-MD5 is correct.  It's highly unlikely the MD5 can be corrupted
  # in a way that also allows corrupt data to pass through.
  #
  # The Rainbows!/Unicorn HTTP servers will populate the HTTP_CONTENT_MD5
  # field in +env+ after env["rack.input"] is fully-consumed.  Clients
  # may also send Content-MD5 as a header and this will still work.
  def content_md5_fail?(env, received) # :nodoc:
    expected = env["HTTP_CONTENT_MD5"] or return false
    expected = expected.strip
    expected == received and return false # success
    r(400, "Content-MD5 mismatch\n" \
           "expected: #{expected}\n" \
           "received: #{received}", env)
  end

end