flexnbd: Add a proxy mode

This lets us proxy connections between NBD clients and servers, resiliently.

tests/acceptance/custom.supp

@@ -0,0 +1,13 @@
+{
+   avoid_glibc_bug_do_lookup
+   Memcheck:Addr8
+   fun:do_lookup_x
+   obj:*
+   fun:_dl_lookup_symbol_x
+}
+{
+   avoid_glibc_bug_check_match
+   Memcheck:Addr8
+   fun:check_match.12149
+}
+

tests/acceptance/environment.rb

@@ -21,6 +21,13 @@ class Environment
     @fake_pid = nil
   end
 
+  def proxy1(port=@port2)
+    @nbd1.proxy(@ip, port)
+  end
+  def proxy2(port=@port1)
+    @nbd2.proxy(@ip, port)
+  end
+
 
   def serve1(*acl)
     @nbd1.serve(@filename1, *acl)

tests/acceptance/flexnbd.rb

@@ -21,7 +21,7 @@ class ValgrindExecutor
   attr_reader :pid
 
   def run( cmd )
-    @pid = fork do exec "valgrind --track-origins=yes #{cmd}" end
+    @pid = fork do exec "valgrind --suppressions=custom.supp --track-origins=yes #{cmd}" end
   end
 end # class ValgrindExecutor
 
@@ -131,7 +131,7 @@ class ValgrindKillingExecutor
 
   def run( cmd )
     @io_r, io_w = IO.pipe
-    @pid = fork do exec( "valgrind --xml=yes --xml-fd=#{io_w.fileno} " + cmd ) end
+    @pid = fork do exec( "valgrind --suppressions=custom.supp --xml=yes --xml-fd=#{io_w.fileno} " + cmd ) end
     launch_watch_thread( @pid, @io_r )
     @pid
   end
@@ -241,6 +241,15 @@ module FlexNBD
         "#{acl.join(' ')}"
     end
 
+    def proxy_cmd( connect_ip, connect_port )
+      "#{bin} proxy "\
+        "--addr #{ip} "\
+        "--port #{port} "\
+        "--conn-addr #{connect_ip} "\
+        "--conn-port #{connect_port} "\
+        "#{@debug}"
+    end
+
 
     def read_cmd( offset, length )
       "#{bin} read "\
@@ -319,6 +328,7 @@ module FlexNBD
         sleep 0.1
       end
 
+
       start_wait_thread( @pid )
       at_exit { kill }
     end
@@ -336,6 +346,31 @@ module FlexNBD
       run_serve_cmd( listen_cmd( file, acl ) )
     end
 
+    def tcp_server_open?
+      # raises if the other side doesn't accept()
+      sock = TCPSocket.new(ip, port) rescue nil
+
+      success = !!sock
+      ( sock.close rescue nil) if sock
+      success
+    end
+
+    def proxy( connect_ip, connect_port )
+      cmd = proxy_cmd( connect_ip, connect_port )
+      debug( cmd )
+
+      @pid = @executor.run( cmd )
+
+      until tcp_server_open?
+        pid, status = Process.wait2(@pid, Process::WNOHANG)
+        raise "server did not start (#{cmd})" if pid
+        sleep 0.1
+      end
+
+      start_wait_thread( @pid )
+      at_exit { kill }
+    end
+
 
     def start_wait_thread( pid )
       @wait_thread = Thread.start do
@@ -521,3 +556,4 @@ module FlexNBD
   end
 
 end
+

tests/acceptance/test_proxy_mode.rb

@@ -0,0 +1,199 @@
+require 'test/unit'
+require 'environment'
+require 'flexnbd/fake_source'
+require 'flexnbd/fake_dest'
+
+class TestProxyMode < Test::Unit::TestCase
+  def setup
+    super
+    @env = Environment.new
+    @env.writefile1( "0" * 16 )
+  end
+
+  def teardown
+    @env.cleanup
+    super
+  end
+
+  def with_proxied_client( override_size = nil )
+    @env.serve1 unless @server_up
+    @env.proxy2 unless @proxy_up
+    @env.nbd2.can_die(0)
+    client = FlexNBD::FakeSource.new(@env.ip, @env.port2, "Couldn't connect to proxy")
+    begin
+
+      result = client.read_hello
+      assert_equal "NBDMAGIC", result[:magic]
+      assert_equal override_size || @env.file1.size, result[:size]
+
+      yield client
+    ensure
+      client.close rescue nil
+    end
+  end
+
+  def test_exits_with_error_when_cannot_connect_to_upstream_on_start
+    assert_raises(RuntimeError) { @env.proxy1 }
+  end
+
+  def test_read_requests_successfully_proxied
+    with_proxied_client do |client|
+      (0..3).each do |n|
+        offset = n * 4096
+        client.write_read_request(offset, 4096, "myhandle")
+        rsp = client.read_response
+
+        assert_equal ::FlexNBD::REPLY_MAGIC, rsp[:magic]
+        assert_equal "myhandle", rsp[:handle]
+        assert_equal 0, rsp[:error]
+
+        orig_data = @env.file1.read(offset, 4096)
+        data = client.read_raw(4096)
+
+        assert_equal 4096, orig_data.size
+        assert_equal 4096, data.size
+
+        assert_equal( orig_data, data, "Returned data does not match" )
+      end
+    end
+  end
+
+  def test_write_requests_successfully_proxied
+    with_proxied_client do |client|
+      (0..3).each do |n|
+        offset = n * 4096
+        client.write(offset, "\xFF" * 4096)
+        rsp = client.read_response
+
+        assert_equal FlexNBD::REPLY_MAGIC, rsp[:magic]
+        assert_equal "myhandle", rsp[:handle]
+        assert_equal 0, rsp[:error]
+
+        data = @env.file1.read(offset, 4096)
+        assert_equal( ( "\xFF" * 4096 ), data, "Data not written" )
+      end
+    end
+  end
+
+  def make_fake_server
+    server = FlexNBD::FakeDest.new(@env.ip, @env.port1)
+    @server_up = true
+
+    # We return a thread here because accept() and connect() both block for us
+    Thread.new do
+      sc = server.accept # just tell the supervisor we're up
+      sc.write_hello
+
+      [ server, sc ]
+    end
+  end
+
+  def test_read_request_retried_when_upstream_dies_partway
+    maker = make_fake_server
+
+    with_proxied_client(4096) do |client|
+      server, sc1 = maker.value
+
+      # Send the read request to the proxy
+      client.write_read_request( 0, 4096 )
+
+      # ensure we're given the read request
+      req1 = sc1.read_request
+      assert_equal ::FlexNBD::REQUEST_MAGIC, req1[:magic]
+      assert_equal ::FlexNBD::REQUEST_READ, req1[:type]
+      assert_equal 0, req1[:from]
+      assert_equal 4096, req1[:len]
+
+      # Kill the server again, now we're sure the read request has been sent once
+      sc1.close
+
+      # We expect the proxy to reconnect without our client doing anything.
+      sc2 = server.accept
+      sc2.write_hello
+
+      # And once reconnected, it should resend an identical request.
+      req2 = sc2.read_request
+      assert_equal req1, req2
+
+      # The reply should be proxied back to the client.
+      sc2.write_reply( req2[:handle] )
+      sc2.write_data( "\xFF" * 4096 )
+
+      # Check it to make sure it's correct
+      rsp = timeout(15) { client.read_response }
+      assert_equal ::FlexNBD::REPLY_MAGIC, rsp[:magic]
+      assert_equal 0, rsp[:error]
+      assert_equal req1[:handle], rsp[:handle]
+
+      data = client.read_raw( 4096 )
+      assert_equal( ("\xFF" * 4096), data, "Wrong data returned" )
+
+      sc2.close
+      server.close
+    end
+
+  end
+
+  def test_write_request_retried_when_upstream_dies_partway
+    maker = make_fake_server
+
+    with_proxied_client(4096) do |client|
+      server, sc1 = maker.value
+
+      # Send the read request to the proxy
+      client.write( 0, ( "\xFF" * 4096 ) )
+
+      # ensure we're given the read request
+      req1 = sc1.read_request
+      assert_equal ::FlexNBD::REQUEST_MAGIC, req1[:magic]
+      assert_equal ::FlexNBD::REQUEST_WRITE, req1[:type]
+      assert_equal 0, req1[:from]
+      assert_equal 4096, req1[:len]
+      data1 = sc1.read_data( 4096 )
+      assert_equal( ( "\xFF" * 4096 ), data1, "Data not proxied successfully" )
+
+      # Kill the server again, now we're sure the read request has been sent once
+      sc1.close
+
+      # We expect the proxy to reconnect without our client doing anything.
+      sc2 = server.accept
+      sc2.write_hello
+
+      # And once reconnected, it should resend an identical request.
+      req2 = sc2.read_request
+      assert_equal req1, req2
+      data2 = sc2.read_data( 4096 )
+      assert_equal data1, data2
+
+      # The reply should be proxied back to the client.
+      sc2.write_reply( req2[:handle] )
+
+      # Check it to make sure it's correct
+      rsp = timeout(15) { client.read_response }
+      assert_equal ::FlexNBD::REPLY_MAGIC, rsp[:magic]
+      assert_equal 0, rsp[:error]
+      assert_equal req1[:handle], rsp[:handle]
+
+      sc2.close
+      server.close
+    end
+  end
+
+  def test_only_one_client_can_connect_to_proxy_at_a_time
+    with_proxied_client do |client|
+
+      c2 = nil
+      assert_raises(Timeout::Error) do
+        timeout(1) do
+          c2 = FlexNBD::FakeSource.new(@env.ip, @env.port2, "Couldn't connect to proxy (2)")
+          c2.read_hello
+        end
+      end
+      c2.close rescue nil if c2
+    end
+
+
+  end
+
+end
+