Replaces sshtunnel with a direct paramiko SSH tunnel

coriolis/providers/replicator.py

@@ -3,15 +3,17 @@
 
 import json
 import os
+import select
 import shutil
+import socket
 import tempfile
+import threading
 import time
 
 from oslo_config import cfg
 from oslo_log import log as logging
 from oslo_utils import units
 import paramiko
-from sshtunnel import SSHTunnelForwarder
 
 from coriolis import exception
 from coriolis.providers import provider_utils
@@ -44,6 +46,95 @@
 CONF.register_opts(replicator_opts, 'replicator')
 
 
+class _SSHTunnel(object):
+    """Local port-forward SSH tunnel backed by paramiko."""
+
+    def __init__(self, ssh_host, ssh_port, ssh_username, ssh_pkey,
+                 ssh_password, remote_bind_address, local_bind_address=None):
+        self._ssh_host = ssh_host
+        self._ssh_port = ssh_port
+        self._ssh_username = ssh_username
+        self._ssh_pkey = ssh_pkey
+        self._ssh_password = ssh_password
+        self._remote_bind_address = remote_bind_address
+        self._requested_local = local_bind_address or ('127.0.0.1', 0)
+        self._transport = None
+        self._listen_sock = None
+        self.local_bind_address = None
+
+    def start(self):
+        client = utils.connect_ssh(
+            self._ssh_host,
+            self._ssh_port,
+            self._ssh_username,
+            pkey=self._ssh_pkey,
+            password=self._ssh_password,
+        )
+        self._transport = client.get_transport()
+
+        self._listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self._listen_sock.bind(self._requested_local)
+        self._listen_sock.listen(5)
+        self.local_bind_address = self._listen_sock.getsockname()
+
+        t = threading.Thread(target=self._accept_loop, daemon=True)
+        t.start()
+
+    def _accept_loop(self):
+        while True:
+            try:
+                sock, _ = self._listen_sock.accept()
+            except OSError:
+                return
+
+            t = threading.Thread(
+                target=self._forward, args=(sock,), daemon=True)
+            t.start()
+
+    def _forward(self, local_sock):
+        try:
+            chan = self._transport.open_channel(
+                'direct-tcpip',
+                self._remote_bind_address,
+                local_sock.getpeername(),
+            )
+        except Exception:
+            local_sock.close()
+            return
+
+        try:
+            while True:
+                r, _, _ = select.select([local_sock, chan], [], [])
+                if local_sock in r:
+                    data = local_sock.recv(4096)
+                    if not data:
+                        break
+                    chan.send(data)
+
+                if chan in r:
+                    data = chan.recv(4096)
+                    if not data:
+                        break
+                    local_sock.send(data)
+        finally:
+            chan.close()
+            local_sock.close()
+
+    def stop(self):
+        if self._listen_sock:
+            try:
+                self._listen_sock.close()
+            except OSError:
+                pass
+
+        if self._transport:
+            try:
+                self._transport.close()
+            except Exception:
+                pass
+
+
 class Client(object):
 
     def __init__(self, ip, port, credentials, ssh_conn_info,
@@ -138,8 +229,9 @@ def _get_ssh_tunnel(self):
             raise exception.CoriolisException(
                 "Either password or pkey is required")
 
-        server = SSHTunnelForwarder(
-            (remote_host, remote_port),
+        server = _SSHTunnel(
+            ssh_host=remote_host,
+            ssh_port=remote_port,
             ssh_username=remote_user,
             ssh_pkey=pkey,
             ssh_password=password,

coriolis/tests/integration/base.py

@@ -245,6 +245,10 @@ class ReplicaIntegrationTestBase(CoriolisIntegrationTestBase):
     _CREATE_MINION_POOLS = False
     _SCSI_DEBUG_SIZE_MB = 16
 
+    # Extra source_environment entries merged into the default transfer's
+    # source_environment.
+    _EXTRA_SOURCE_ENVIRONMENT = {}
+
     @classmethod
     def setUpClass(cls):
         harness._IntegrationHarness.get().imp_provider.check_prerequisites()
@@ -307,6 +311,7 @@ def setUp(self):
             source_environment={
                 "instance_block_devices": {
                     self._instance_name: [self._src_device]},
+                **self._EXTRA_SOURCE_ENVIRONMENT,
             },
         )
         # Safety-net cleanup for destination devices allocated by the provider.

coriolis/tests/integration/test_provider/exp.py

@@ -78,7 +78,9 @@ def _make_replicator(self, conn_info, event_mgr, volumes_info, repl_state):
         """Build a Replicator that connects via SSH to *conn_info*.
 
         *conn_info* must contain ``ip``, ``port``, ``username``, and
-        ``pkey_path`` keys.
+        ``pkey_path`` keys. An optional ``use_tunnel`` key forces the
+        replicator client to connect through an SSH tunnel instead of
+        directly to the replicator's TCP port.
         """
         pkey = paramiko.RSAKey.from_private_key_file(conn_info["pkey_path"])
         repl_conn_info = {
@@ -88,7 +90,8 @@ def _make_replicator(self, conn_info, event_mgr, volumes_info, repl_state):
             "pkey": pkey,
         }
         return replicator_module.Replicator(
-            repl_conn_info, event_mgr, volumes_info, repl_state)
+            repl_conn_info, event_mgr, volumes_info, repl_state,
+            use_tunnel=conn_info.get("use_tunnel", False))
 
     # BaseProvider / BaseEndpointProvider
 
@@ -114,6 +117,7 @@ def get_source_environment_schema(self):
             "type": "object",
             "properties": {
                 "instance_block_devices": {"type": "object"},
+                "use_tunnel": {"type": "boolean"},
             },
         }
 
@@ -254,6 +258,7 @@ def deploy_replica_source_resources(
                 "port": 22,
                 "username": "root",
                 "pkey_path": pkey_path,
+                "use_tunnel": source_environment.get("use_tunnel", False),
             }
             replicator = self._make_replicator(
                 src_conn_info, self._event_manager(), [], None)

coriolis/tests/integration/transfers/test_transfer.py

@@ -25,6 +25,7 @@
 from coriolis import data_transfer
 from coriolis.db import api as db_api
 from coriolis.providers import backup_writers
+from coriolis.providers import replicator as replicator_module
 from coriolis.tests.integration import base
 from coriolis.tests.integration import utils as test_utils
 
@@ -398,3 +399,35 @@ def test_transfer(self):
         super().test_transfer()
         self.assertPoolAllocated(self._pool_id)
         self.assertMachinesAvailable(self._pool_id)
+
+
+class ReplicaTransferViaSSHTunnelTest(base.ReplicaIntegrationTestBase):
+    """Transfer tests using an SSH tunneled replicator client."""
+
+    _EXTRA_SOURCE_ENVIRONMENT = {"use_tunnel": True}
+
+    def test_transfer_via_ssh_tunnel(self):
+        tunnel_starts = []
+        original_get_ssh_tunnel = replicator_module.Client._get_ssh_tunnel
+
+        def _spy_get_ssh_tunnel(client_self):
+            tunnel = original_get_ssh_tunnel(client_self)
+            tunnel.start = mock.Mock(wraps=tunnel.start)
+            tunnel_starts.append(tunnel.start)
+            return tunnel
+
+        with mock.patch.object(
+                replicator_module.Client, "_get_ssh_tunnel",
+                _spy_get_ssh_tunnel):
+            self._execute_and_wait(self._transfer.id)
+
+        self.assertTrue(tunnel_starts, "SSH tunnel was never constructed")
+        self.assertTrue(
+            any(t.called for t in tunnel_starts),
+            "SSH tunnel was constructed but never started")
+
+        if self._harness.uses_core_test_import_provider():
+            self.assertTrue(
+                test_utils.devices_match(self._src_device, self._dst_device),
+                "Devices do not match after transfer via SSH tunnel",
+            )

coriolis/tests/providers/test_replicator.py

@@ -15,6 +15,130 @@
 from coriolis.tests import testutils
 
 
+class SSHTunnelTestCase(test_base.CoriolisBaseTestCase):
+    """Test suite for the Coriolis _SSHTunnel class."""
+
+    def setUp(self):
+        super(SSHTunnelTestCase, self).setUp()
+        self.ssh_host = mock.sentinel.ssh_host
+        self.ssh_port = mock.sentinel.ssh_port
+        self.ssh_username = mock.sentinel.ssh_username
+        self.ssh_pkey = mock.sentinel.ssh_pkey
+        self.ssh_password = mock.sentinel.ssh_password
+        self.remote_bind_address = ("127.0.0.1", 1234)
+        self.tunnel = replicator_module._SSHTunnel(
+            self.ssh_host, self.ssh_port, self.ssh_username, self.ssh_pkey,
+            self.ssh_password, self.remote_bind_address)
+
+    @mock.patch.object(replicator_module.threading, 'Thread')
+    @mock.patch.object(replicator_module.socket, 'socket')
+    @mock.patch.object(replicator_module.utils, 'connect_ssh')
+    def test_start(self, mock_connect_ssh, mock_socket, mock_thread):
+        mock_client = mock_connect_ssh.return_value
+        mock_sock = mock_socket.return_value
+        mock_sock.getsockname.return_value = ("127.0.0.1", 4321)
+
+        self.tunnel.start()
+
+        mock_connect_ssh.assert_called_once_with(
+            self.ssh_host, self.ssh_port, self.ssh_username,
+            pkey=self.ssh_pkey, password=self.ssh_password)
+        self.assertEqual(
+            mock_client.get_transport.return_value, self.tunnel._transport)
+
+        mock_socket.assert_called_once_with(
+            replicator_module.socket.AF_INET,
+            replicator_module.socket.SOCK_STREAM)
+        mock_sock.setsockopt.assert_called_once_with(
+            replicator_module.socket.SOL_SOCKET,
+            replicator_module.socket.SO_REUSEADDR, 1)
+        mock_sock.bind.assert_called_once_with(
+            self.tunnel._requested_local)
+        mock_sock.listen.assert_called_once_with(5)
+        self.assertEqual(mock_sock, self.tunnel._listen_sock)
+        self.assertEqual(("127.0.0.1", 4321), self.tunnel.local_bind_address)
+
+        mock_thread.assert_called_once_with(
+            target=self.tunnel._accept_loop, daemon=True)
+        mock_thread.return_value.start.assert_called_once_with()
+
+    @mock.patch.object(replicator_module.threading, 'Thread')
+    def test_accept_loop(self, mock_thread):
+        mock_sock = mock.sentinel.sock
+        self.tunnel._listen_sock = mock.MagicMock()
+        self.tunnel._listen_sock.accept.side_effect = [
+            (mock_sock, mock.sentinel.addr), OSError]
+
+        self.tunnel._accept_loop()
+
+        self.assertEqual(2, self.tunnel._listen_sock.accept.call_count)
+        mock_thread.assert_called_once_with(
+            target=self.tunnel._forward, args=(mock_sock,), daemon=True)
+        mock_thread.return_value.start.assert_called_once_with()
+
+    def test_forward_open_channel_fails(self):
+        local_sock = mock.MagicMock()
+        self.tunnel._transport = mock.MagicMock()
+        self.tunnel._transport.open_channel.side_effect = Exception(
+            "connection refused")
+
+        self.tunnel._forward(local_sock)
+
+        local_sock.close.assert_called_once_with()
+
+    @mock.patch.object(replicator_module.select, 'select')
+    def test_forward(self, mock_select):
+        local_sock = mock.MagicMock()
+        chan = mock.MagicMock()
+        self.tunnel._transport = mock.MagicMock()
+        self.tunnel._transport.open_channel.return_value = chan
+
+        local_sock.recv.return_value = b"request"
+        chan.recv.side_effect = [b"response", b""]
+        mock_select.side_effect = [
+            ([local_sock], [], []),
+            ([chan], [], []),
+            ([chan], [], []),
+        ]
+
+        self.tunnel._forward(local_sock)
+
+        self.tunnel._transport.open_channel.assert_called_once_with(
+            'direct-tcpip', self.tunnel._remote_bind_address,
+            local_sock.getpeername.return_value)
+        chan.send.assert_called_once_with(b"request")
+        local_sock.send.assert_called_once_with(b"response")
+        chan.close.assert_called_once_with()
+        local_sock.close.assert_called_once_with()
+
+    @mock.patch.object(replicator_module.select, 'select')
+    def test_forward_local_sock_closed(self, mock_select):
+        local_sock = mock.MagicMock()
+        chan = mock.MagicMock()
+        self.tunnel._transport = mock.MagicMock()
+        self.tunnel._transport.open_channel.return_value = chan
+
+        local_sock.recv.return_value = b""
+        mock_select.return_value = ([local_sock], [], [])
+
+        self.tunnel._forward(local_sock)
+
+        chan.close.assert_called_once_with()
+        local_sock.close.assert_called_once_with()
+
+    def test_stop(self):
+        self.tunnel._listen_sock = mock.MagicMock()
+        self.tunnel._listen_sock.close.side_effect = OSError
+        self.tunnel._transport = mock.MagicMock()
+        self.tunnel._transport.close.side_effect = Exception(
+            "boom goes the dynamite")
+
+        self.tunnel.stop()
+
+        self.tunnel._listen_sock.close.assert_called_once_with()
+        self.tunnel._transport.close.assert_called_once_with()
+
+
 class ClientTestCase(test_base.CoriolisBaseTestCase):
     """Test suite for the Coriolis Client class."""
 
@@ -145,20 +269,21 @@ def test_test_connection_with_tunnel_with_exception(
         mock_setup_tunnel.assert_called_once()
         mock_tunnel.stop.assert_called_once()
 
-    @mock.patch.object(replicator_module, 'SSHTunnelForwarder')
-    def test__get_ssh_tunnel(self, mock_SSHTunnelForwarder):
+    @mock.patch.object(replicator_module, '_SSHTunnel')
+    def test__get_ssh_tunnel(self, mock_SSHTunnel):
         result = self.client._get_ssh_tunnel()
 
-        mock_SSHTunnelForwarder.assert_called_once_with((
-            self.ssh_conn_info["hostname"], self.ssh_conn_info["port"]),
+        mock_SSHTunnel.assert_called_once_with(
+            ssh_host=self.ssh_conn_info["hostname"],
+            ssh_port=self.ssh_conn_info["port"],
             ssh_username=self.ssh_conn_info["username"],
             ssh_pkey=self.ssh_conn_info["pkey"],
             ssh_password=self.ssh_conn_info["password"],
             remote_bind_address=("127.0.0.1", self.port),
             local_bind_address=("127.0.0.1", 0)
         )
 
-        self.assertEqual(result, mock_SSHTunnelForwarder.return_value)
+        self.assertEqual(result, mock_SSHTunnel.return_value)
 
     def test__get_ssh_tunnel_no_credentials(self):
         self.client._ssh_conn_info = {

requirements.txt

@@ -44,7 +44,6 @@ strict-rfc3339
 sqlalchemy<2.0.0
 taskflow
 webob
-sshtunnel
 requests-unixsocket
 # Cherrypy wsgi, also used by Ironic.
 # Leveraged by coriolis-api when called directly.