#!/usr/bin/ruby

require 'test/unit'
require 'flexnbd'
require 'test_file_writer'

class Environment
  attr_reader( :blocksize, :filename1, :filename2, :ip,
               :port1, :port2, :nbd1, :nbd2, :file1, :file2 )

  def initialize
    @blocksize = 1024
    @filename1 = "/tmp/.flexnbd.test.#{$$}.#{Time.now.to_i}.1"
    @filename2 = "/tmp/.flexnbd.test.#{$$}.#{Time.now.to_i}.2"
    @ip = "127.0.0.1"
    @available_ports = [*40000..41000] - listening_ports
    @port1 = @available_ports.shift
    @port2 = @available_ports.shift
    @nbd1 = FlexNBD.new("../build/flexnbd", @ip, @port1)
    @nbd2 = FlexNBD.new("../build/flexnbd", @ip, @port2)

    @fake_pid = nil
  end


  def serve1(*acl)
    @nbd1.serve(@filename1, *acl)
  end

  def serve2(*acl)
    @nbd2.serve(@filename2, *acl)
  end


  def listen1( *acl )
    @nbd1.listen( @filename1, *(acl.empty? ? @acl1: acl) )
  end

  def listen2( *acl )
    @nbd2.listen( @filename2, *acl )
  end


  def acl1( *acl )
    @nbd1.acl( *acl )
  end

  def acl2( *acl )
    @nbd2.acl( *acl )
  end


  def status1
    @nbd1.status.first
  end

  def status2
    @nbd2.status.first
  end



  def mirror12
    @nbd1.mirror( @nbd2.ip, @nbd2.port )
  end

  def mirror12_unchecked
    @nbd1.mirror_unchecked( @nbd2.ip, @nbd2.port, nil, nil, 10 )
  end


  def writefile1(data)
    @file1 = TestFileWriter.new(@filename1, @blocksize).write(data)
  end

  def writefile2(data)
    @file2 = TestFileWriter.new(@filename2, @blocksize).write(data)
  end



  def listening_ports
    `netstat -ltn`.
      split("\n").
      map { |x| x.split(/\s+/) }[2..-1].
      map { |l| l[3].split(":")[-1].to_i }
  end


  def cleanup
    if @fake_pid
      begin
        Process.waitpid2( @fake_pid )
      rescue Errno::ESRCH
      end
    end


    @nbd1.kill
    @nbd2.kill

    [@filename1, @filename2].each do |f|
      File.unlink(f) if File.exists?(f)
    end
  end


  def run_fake( name, addr, port )
    fakedir = File.join( File.dirname( __FILE__ ), "fakes" )
    fake = Dir[File.join( fakedir, name ) + "*"].sort.find { |fn|
      File.executable?( fn )
    }

    raise "no fake executable" unless fake
    raise "no addr" unless addr
    raise "no port" unless port
    @fake_pid = fork do
      exec fake + " " + addr.to_s + " " + port.to_s
    end
    sleep(0.5)
  end


  def fake_reports_success
    _,status = Process.waitpid2( @fake_pid )
    @fake_pid = nil
    status.success?
  end


end # class Environment


class NBDScenarios < Test::Unit::TestCase
  def setup
    @env = Environment.new
  end

  def teardown
    @env.nbd1.can_die(0)
    @env.nbd2.can_die(0)
    @env.cleanup
  end


  def test_read1
    @env.writefile1("f"*64)
    @env.serve1

    [0, 12, 63].each do |num|

      assert_equal(
        @env.nbd1.read(num*@env.blocksize, @env.blocksize),
        @env.file1.read(num*@env.blocksize, @env.blocksize)
      )
    end

    [124, 1200, 10028, 25488].each do |num|
      assert_equal(@env.nbd1.read(num, 4), @env.file1.read(num, 4))
    end
  end

  # Check that we're not
  #
  def test_writeread1
    @env.writefile1("0"*64)
    @env.serve1

    [0, 12, 63].each do |num|
      data = "X"*@env.blocksize
      @env.nbd1.write(num*@env.blocksize, data)
      assert_equal(data, @env.file1.read(num*@env.blocksize, data.size))
      assert_equal(data, @env.nbd1.read(num*@env.blocksize, data.size))
    end
  end

  # Check that we're not overstepping or understepping where our writes end
  # up.
  #
  def test_writeread2
    @env.writefile1("0"*1024)
    @env.serve1

    d0 = "\0"*@env.blocksize
    d1 = "X"*@env.blocksize
    (0..63).each do |num|
      @env.nbd1.write(num*@env.blocksize*2, d1)
    end
    (0..63).each do |num|
      assert_equal(d0, @env.nbd1.read(((2*num)+1)*@env.blocksize, d0.size))
    end
  end


  def test_mirror
    @env.writefile1( "f"*4 )
    @env.serve1

    @env.writefile2( "0"*4 )
    @env.listen2

    @env.nbd1.can_die
    stdout, stderr = @env.mirror12

    @env.nbd1.join

    assert_equal(@env.file1.read_original( 0, @env.blocksize ),
                 @env.file2.read( 0, @env.blocksize ) )
    assert @env.status2['has_control'], "destination didn't take control"
  end

end


class NBDConnectSourceFailureScenarios < Test::Unit::TestCase
  def setup
    @env = Environment.new
    @env.writefile1( "f" * 4 )
    @env.serve1

  end

  def teardown
    @env.nbd1.can_die(0)
    @env.cleanup
  end


  def test_failure_to_connect_reported_in_mirror_cmd_response
    stdout, stderr = @env.mirror12_unchecked
    assert_match( /failed to connect/, stderr )
  end


  def test_destination_hangs_after_connect_reports_error_at_source
    run_fake( "dest/hang_after_connect" )

    stdout, stderr = @env.mirror12_unchecked
    assert_match( /Remote server failed to respond/, stderr )
    assert_success
  end


  def test_destination_rejects_connection_reports_error_at_source
    run_fake( "dest/reject_acl" )

    stdout, stderr = @env.mirror12_unchecked
    assert_match /Mirror was rejected/, stderr
    assert_success
  end

  def test_wrong_size_causes_disconnect
    run_fake( "dest/hello_wrong_size" )
    stdout, stderr = @env.mirror12_unchecked
    assert_match /Remote size does not match local size/, stderr
    assert_success
  end


  def test_wrong_magic_causes_disconnect
    run_fake( "dest/hello_wrong_magic" )
    stdout, stderr = @env.mirror12_unchecked
    assert_match /Mirror was rejected/, stderr
    assert_success "dest/hello_wrong_magic fake failed"
  end


  def test_disconnect_after_hello_causes_retry
    run_fake( "dest/close_after_hello" )
    stdout, stderr = @env.mirror12_unchecked
    assert_match( /Mirror started/, stdout )

    assert_success
  end


  def test_write_times_out_causes_retry
    run_fake( "dest/hang_after_write" )
    stdout, stderr = @env.mirror12_unchecked

    assert_success
  end


  def test_rejected_write_causes_retry
    run_fake( "dest/error_on_write" )
    stdout, stderr = @env.mirror12_unchecked
    assert_success
  end


  private
  def run_fake(name)
    @env.run_fake( name, @env.ip, @env.port2 )
  end

  def assert_success( msg=nil )
    assert @env.fake_reports_success, msg || "Fake failed"
  end


end


class NBDConnectDestFailureScenarios < Test::Unit::TestCase

  def setup
    @env = Environment.new
    @env.writefile1( "0" * 4 )
    @env.listen1
  end

  def teardown
    @env.cleanup
  end


  def test_hello_blocked_by_disconnect_causes_error_not_fatal
    run_fake( "source/close_after_connect" )
    assert_no_control
  end


  def test_hello_goes_astray_causes_timeout_error
    run_fake( "source/hang_after_hello" )
    assert_no_control
  end


  def test_disconnect_after_hello_causes_error_not_fatal
    run_fake( "source/close_after_hello" )
    assert_no_control
  end


  def test_double_connect_during_hello
    run_fake( "source/connect_during_hello" )
  end


  def test_acl_rejection
    @env.acl1("127.0.0.1")
    run_fake( "source/connect_from_banned_ip")
  end


  def test_bad_write
    run_fake( "source/write_out_of_range" )
  end


  private
  def run_fake( name )
    @env.run_fake( name, @env.ip, @env.port1 )
    assert @env.fake_reports_success, "#{name} failed."
  end

  def assert_no_control
    status, stderr = @env.status1
    assert !status['has_control'],  "Thought it had control"
  end


end # class NBDConnectDestFailureScenarios