Factor common code out of fake destinations

tests/fakes/dest/close_after_hello.rb

@@ -6,40 +6,18 @@
 # command-line has gone away, and can't feed an error back to the
 # user, we have to keep trying.
 
+require 'flexnbd/fake_dest'
+include FlexNBD::FakeDest
+
 addr, port = *ARGV
 
-require 'socket'
-require 'timeout'
-
-sock = TCPServer.new( addr, port )
-client_sock = nil
-
-begin
-  Timeout.timeout(2) do
-    client_sock = sock.accept
-  end
-rescue Timeout::Error
-  $stderr.puts "Timed out waiting for a connection"
-  exit 1
-end
-
-
-client_sock.write( "NBDMAGIC" )
-client_sock.write( "\x00\x00\x42\x02\x81\x86\x12\x53" )
-client_sock.write( "\x00\x00\x00\x00\x00\x00\x10\x00" )
-client_sock.write( "\x00" * 128 )
 
+sock = serve( addr, port )
+client_sock = accept( sock, "Timed out waiting for a connection" )
+write_hello( client_sock )
 client_sock.close
 
-new_sock = nil
+new_sock = accept( sock, "Timed out waiting for a reconnection" )
-begin
-  Timeout.timeout(60) do
-    new_sock = sock.accept
-  end
-rescue Timeout::Error
-  $stderr.puts "Timed out waiting for a reconnection"
-  exit 1
-end
 
 new_sock.close
 sock.close

tests/fakes/dest/hang_after_connect.rb

@@ -7,24 +7,13 @@
 # This allows the test runner to check that the command-line sees the
 # right error message after the timeout time.
 
-require 'socket'
+require 'flexnbd/fake_dest'
-require 'timeout'
+include FlexNBD::FakeDest
-require 'flexnbd/constants'
 
 addr, port = *ARGV
 
-serve_sock = TCPServer.new( addr, port )
+serve_sock = serve( addr, port )
-client_sock = nil
+client_sock = accept( serve_sock, "Client didn't make a connection" )
-
-begin
-  # A failure here means a more serious issue with flexnbd
-  Timeout.timeout( 2 ) do
-    client_sock = serve_sock.accept
-  end
-rescue Timeout::Error
-  $stderr.puts "Client didn't make connection"
-  exit 1
-end
 
 # Sleep for one second past the timeout (a bit of slop in case ruby
 # doesn't launch things quickly)

tests/fakes/dest/hello_wrong_magic.rb

@@ -3,43 +3,22 @@
 # Simulate a destination which sends the wrong magic.
 # We expect the sender to disconnect and reconnect.
 
-addr, port = *ARGV
+require 'flexnbd/fake_dest'
-
+include FlexNBD::FakeDest
-require 'socket'
-require 'timeout'
-require 'flexnbd/constants'
-
-sock = TCPServer.new( addr, port )
-client_sock = nil
-
-begin
-  Timeout.timeout(2) do
-    client_sock = sock.accept
-  end
-rescue Timeout::Error
-  $stderr.puts "Timed out waiting for a connection"
-  exit 1
-end
 
+sock = serve( *ARGV )
+client_sock = accept( sock, "Timed out waiting for a connection" )
 
+# Launch a second thread so that we can spot the reconnection attempt
+# as soon as it happens, or alternatively die a flaming death on
+# timeout.
 t = Thread.new do
-  begin
+  client_sock2 = accept( sock, "Timed out waiting for a reconnection",
-    Timeout.timeout(FlexNBD::MS_RETRY_DELAY_SECS + 1) do
+                         FlexNBD::MS_RETRY_DELAY_SECS + 1 )
-      client_sock2 = sock.accept
+  client_sock2.close
-    end
-  rescue Timeout::Error
-    $stderr.puts "Timed out waiting for a reconnection"
-    exit 1
-  end
 end
 
-client_sock.write( "NBDMAGIC" )
+write_hello( client_sock, :magic => :wrong )
-# We're off in the last byte, should be \x53
-client_sock.write( "\x00\x00\x42\x02\x81\x86\x12\x52" )
-# 4096 is the right size; this is defined in nbd_scenarios
-client_sock.write( "\x00\x00\x00\x00\x00\x00\x10\x00" )
-client_sock.write( "\x00" * 128 )
-
 
 t.join
 

tests/fakes/dest/hello_wrong_size.rb

@@ -4,40 +4,20 @@
 # a valid NBD hello with a random size, then check that we have see an
 # EOF on read.
 
-addr, port = *ARGV
+require 'flexnbd/fake_dest'
+include FlexNBD::FakeDest
 
-require 'socket'
+sock = serve( *ARGV )
-require 'timeout'
+client_sock = accept( sock, "Timed out waiting for a connection" )
-require 'flexnbd/constants'
 
-sock = TCPServer.new( addr, port )
-client_sock = nil
-
-begin
-  Timeout.timeout(2) do
-    client_sock = sock.accept
-  end
-rescue Timeout::Error
-  $stderr.puts "Timed out waiting for a connection"
-  exit 1
-end
 
 t = Thread.new do
-  begin
+  client_sock2 = accept( sock, "Timed out waiting for a reconnection",
-    Timeout.timeout(FlexNBD::MS_RETRY_DELAY_SECS + 1) do
+                         FlexNBD::MS_RETRY_DELAY_SECS + 1 )
-      client_sock2 = sock.accept
+  client_sock2.close
-    end
-  rescue Timeout::Error
-    $stderr.puts "Timed out waiting for a reconnection"
-    exit 1
-  end
 end
 
-client_sock.write( "NBDMAGIC" )
+write_hello( client_sock, :size => :wrong )
-client_sock.write( "\x00\x00\x42\x02\x81\x86\x12\x53" )
-8.times do client_sock.write rand(256).chr end
-client_sock.write( "\x00" * 128 )
-
 
 t.join
 

tests/fakes/dest/reject_acl.rb

@@ -2,20 +2,11 @@
 
 # Accept a connection, then immediately close it.  This simulates an ACL rejection.
 
-addr, port = *ARGV
+require 'flexnbd/fake_dest'
-require 'socket'
+include FlexNBD::FakeDest
-require 'timeout'
 
-serve_sock = TCPServer.open( addr, port )
+serve_sock = serve( *ARGV )
-
+accept( serve_sock, "Timed out waiting for a connection" ).close
-begin
-  Timeout.timeout( 2 ) do
-    serve_sock.accept.close
-  end
-rescue Timeout::Error
-  $stderr.puts "Timed out waiting for a connection"
-  exit 1
-end
 
 serve_sock.close
 

tests/flexnbd/fake_dest.rb

@@ -0,0 +1,51 @@
+# encoding: utf-8
+
+require 'socket'
+require 'timeout'
+
+require 'flexnbd/constants'
+
+module FlexNBD
+  module FakeDest
+
+    def serve( addr, port )
+      TCPServer.new( addr, port )
+    end
+
+
+    def accept( sock, err_msg, timeout=2 )
+      client_sock = nil
+
+      begin
+        Timeout.timeout(timeout) do
+          client_sock = sock.accept
+        end
+      rescue Timeout::Error
+        $stderr.puts err_msg
+        exit 1
+      end
+
+      client_sock
+    end
+
+
+    def write_hello( client_sock, opts={} )
+      client_sock.write( "NBDMAGIC" )
+
+      if opts[:magic] == :wrong
+        client_sock.write( "\x00\x00\x42\x02\x81\x86\x12\x52" )
+      else
+        client_sock.write( "\x00\x00\x42\x02\x81\x86\x12\x53" )
+      end
+
+      if opts[:size] == :wrong
+        8.times do client_sock.write rand(256).chr end
+      else
+        client_sock.write( "\x00\x00\x00\x00\x00\x00\x10\x00" )
+      end
+
+      client_sock.write( "\x00" * 128 )
+    end
+
+  end
+end # module FlexNBD