diff --git a/storages/test_ssh_pool.py b/storages/test_ssh_pool.py new file mode 100644 index 00000000..595cb5ff --- /dev/null +++ b/storages/test_ssh_pool.py @@ -0,0 +1,141 @@ +"""Chunks to an SSH compute reuse this process's SSH connection: each request +takes it for itself and opens its own SFTP session. A connection is not +reused after an error, when the host dropped it, when it is too old, or by +another login; an idle one is closed, and a forked worker leaves the +parent's connections to the parent.""" + +import os +import threading +import time +import unittest +from unittest.mock import MagicMock, patch + +from storages import upload +from storages.upload import _ssh_pool, _UploadTarget +from vrtManager.connection import CONN_SSH + + +def client(): + return MagicMock(**{"get_transport.return_value.is_active.return_value": True}) + + +KEY = ("192.0.2.1", 22, "root", "secret") + + +def target(host="192.0.2.1", login="root", passwd="secret"): + return _UploadTarget(MagicMock(conn=CONN_SSH, host=host, login=login, passwd=passwd), "/pool", "x.iso", "id") + + +class SSHPoolTestCase(unittest.TestCase): + def setUp(self): + self.clients = [] + patcher = patch("paramiko.SSHClient", side_effect=lambda: self.clients.append(client()) or self.clients[-1]) + patcher.start() + self.addCleanup(patcher.stop) + self.addCleanup(_ssh_pool.close_all) + + def request(self, **kwargs): + with target(**kwargs) as t: + return t.ssh + + def test_chunks_and_uploads_share_one_connection(self): + first = self.request() + self.assertIs(self.request(), first) + self.assertEqual(len(self.clients), 1) + self.assertEqual(first.open_sftp.call_count, 2) + first.close.assert_not_called() + + def test_another_login_port_or_password_has_its_own_connection(self): + for kwargs in ({}, {"login": "admin"}, {"host": "192.0.2.1:2222"}, {"passwd": "changed"}): + self.request(**kwargs) + self.assertEqual(len(self.clients), 4) + self.assertEqual(self.clients[2].connect.call_args.kwargs["port"], 2222) + + def test_two_requests_at_once_do_not_share_and_the_surplus_is_closed(self): + with target() as a, target() as b: + self.assertIsNot(a.ssh, b.ssh) + first, second = self.clients # b is given back first and kept + first.close.assert_called_once() + second.close.assert_not_called() + self.assertIs(self.request(), second) + + def test_a_connection_is_not_reused_after_an_error(self): + with self.assertRaises(OSError): + with target(): + raise OSError("write failed") + self.clients[0].close.assert_called_once() + self.request() + self.assertEqual(len(self.clients), 2) + + def test_a_dropped_connection_is_replaced(self): + stale = self.request() + stale.get_transport.return_value.is_active.return_value = False + self.assertIsNot(self.request(), stale) + stale.close.assert_called_once() + + def test_a_connection_whose_sftp_fails_is_replaced(self): + stale = self.request() + stale.open_sftp.side_effect = EOFError("connection reset") + self.assertIsNot(self.request(), stale) + stale.close.assert_called_once() + + def test_an_sftp_start_that_never_answers_is_ended(self): + stale = self.request() + closed = threading.Event() + + def hang(): # paramiko's wait for the subsystem, woken by a close + closed.wait(5) + raise EOFError("closed") + + stale.open_sftp.side_effect = hang + stale.close.side_effect = closed.set + start = time.monotonic() + with patch("storages.upload.SSH_CONNECT_TIMEOUT_SECONDS", 0.05): + self.assertIsNot(self.request(), stale) + self.assertLess(time.monotonic() - start, 1) + + def test_a_connection_without_an_idle_timer_is_not_kept(self): + idle = client() + with patch("threading.Timer.start", side_effect=RuntimeError("can't start new thread")): + _ssh_pool.give(KEY, idle, time.monotonic()) + idle.close.assert_called_once() + self.assertEqual(_ssh_pool._idle, {}) + + def test_a_connection_is_not_reused_after_its_lifetime(self): + stale = self.request() + with patch("storages.upload.time.monotonic", return_value=time.monotonic() + upload.SSH_MAX_AGE_SECONDS): + self.assertIsNot(self.request(), stale) + stale.close.assert_called_once() + + def test_an_idle_connection_is_closed(self): + with patch("storages.upload.SSH_IDLE_SECONDS", 0.01): + idle = self.request() + time.sleep(0.2) + idle.close.assert_called_once() + self.assertIsNot(self.request(), idle) + + def test_an_old_timer_does_not_close_a_connection_given_back_since(self): + first = self.request() + old_timer = _ssh_pool._idle[KEY][2] + self.request() # taken and given back: a new timer + _ssh_pool._expire(KEY, old_timer) + first.close.assert_not_called() + self.assertIs(self.request(), first) + + def test_a_forked_child_drops_the_parents_connections(self): + parent = self.request() + with _ssh_pool._lock: # held while forking + pid = os.fork() + if pid == 0: + ok = _ssh_pool._idle == {} and parent.get_transport.return_value.atfork.called + with _ssh_pool._lock: + pass + os._exit(0 if ok and not parent.close.called else 1) + _, status = os.waitpid(pid, 0) + self.assertEqual(os.waitstatus_to_exitcode(status), 0) + parent.get_transport.return_value.atfork.assert_not_called() + self.assertIs(self.request(), parent) + + +if __name__ == "__main__": + unittest.main() diff --git a/storages/tests.py b/storages/tests.py index d5bc9215..b4981023 100644 --- a/storages/tests.py +++ b/storages/tests.py @@ -18,6 +18,7 @@ from storages.upload import ( _get_upload_lock_directory, _is_session_completed, _sftp_replace_file, + _ssh_pool, ) from vrtManager.connection import CONN_SOCKET, CONN_SSH, CONN_TCP, CONN_TLS @@ -117,6 +118,7 @@ class StorageUploadSecurityTests(TestCase): self.mock_ssh_conn.host = "127.0.0.1" self.mock_ssh_conn.login = "test" self.mock_ssh_conn.passwd = "test" + self.addCleanup(_ssh_pool.close_all) def tearDown(self): shutil.rmtree(self.temp_dir, ignore_errors=True) @@ -635,6 +637,9 @@ class StorageUploadSecurityTests(TestCase): upload_id=upload_id, ) self.assertEqual(remote_files[final_remote], b"DEBIAN_CHUNK0_DEBIAN_CHUNK1") + # one connection for the three requests, an SFTP session for each + mock_ssh.connect.assert_called_once() + self.assertEqual(mock_ssh.open_sftp.call_count, 3) def test_handle_uploaded_file_ssh_traversal_check(self): dummy_chunk = SimpleUploadedFile("test.iso", b"data") diff --git a/storages/upload.py b/storages/upload.py index 0454465e..34b5cb8b 100644 --- a/storages/upload.py +++ b/storages/upload.py @@ -1,5 +1,6 @@ """Chunked ISO upload operations shared by the storage view.""" +import atexit import fcntl import json import logging @@ -8,6 +9,7 @@ import posixpath import re import stat import tempfile +import threading import time from xml.sax.saxutils import escape @@ -20,6 +22,86 @@ logger = logging.getLogger(__name__) LOCK_TIMEOUT_SECONDS = 15 SSH_CONNECT_TIMEOUT_SECONDS = 10 SSH_IO_TIMEOUT_SECONDS = 30 +SSH_IDLE_SECONDS = 60 +SSH_MAX_AGE_SECONDS = 600 + + +class _SSHPool: + """SSH connections to computes, kept for the next chunk of this process + (one HTTP request per chunk; a new connection costs ~200 ms). A request + takes a connection for itself and opens its own SFTP session on it. An + idle connection is closed after SSH_IDLE_SECONDS, and none is reused + once it is SSH_MAX_AGE_SECONDS old, so host keys and logins are checked + again.""" + + def __init__(self): + self._lock = threading.Lock() + self._idle = {} # key -> (client, opened, timer) + + @staticmethod + def _usable(client, opened): + transport = client.get_transport() + return transport is not None and transport.is_active() and time.monotonic() - opened < SSH_MAX_AGE_SECONDS + + def take(self, key): + """An idle connection and when it was opened, or None""" + with self._lock: + entry = self._idle.pop(key, None) + if entry is None: + return None + client, opened, timer = entry + timer.cancel() + if self._usable(client, opened): + return client, opened + client.close() + return None + + def give(self, key, client, opened): + if not self._usable(client, opened): + client.close() + return + timer = threading.Timer(SSH_IDLE_SECONDS, lambda: self._expire(key, timer)) + timer.daemon = True + with self._lock: + kept = key not in self._idle + if kept: + try: + timer.start() + except RuntimeError: # no thread for the timer: not kept + kept = False + else: + self._idle[key] = (client, opened, timer) + if not kept: + client.close() + + def _expire(self, key, timer): + with self._lock: + entry = self._idle.get(key) + if entry is None or entry[2] is not timer: # taken since, perhaps given back + return + del self._idle[key] + entry[0].close() + + def close_all(self): + with self._lock: + entries, self._idle = list(self._idle.values()), {} + for client, _opened, timer in entries: + timer.cancel() + client.close() + + def forget(self): + """For a forked child: drop the parent's connections without ending + their sessions; the parent keeps using them""" + entries, self._lock, self._idle = list(self._idle.values()), threading.Lock(), {} + for client, _opened, _timer in entries: + transport = client.get_transport() + if transport is not None: + transport.atfork() + + +_ssh_pool = _SSHPool() +os.register_at_fork(after_in_child=_ssh_pool.forget) +atexit.register(_ssh_pool.close_all) def _fsync_dir(dir_path): @@ -348,32 +430,71 @@ class _UploadTarget: def __enter__(self): if self.remote: hostname, separator, port = self.conn.host.partition(":") - self.ssh = paramiko.SSHClient() + self.key = (hostname, int(port) if separator else 22, self.conn.login, self.conn.passwd) + taken = _ssh_pool.take(self.key) + if taken is not None: + self.ssh, self.opened = taken + try: + self._open_sftp() + return self + except Exception: # a connection the host dropped: open a new one + self.ssh.close() + self.ssh, self.opened = self._connect(), time.monotonic() try: - self.ssh.load_system_host_keys() - self.ssh.set_missing_host_key_policy(paramiko.RejectPolicy()) - self.ssh.connect( - hostname=hostname, port=int(port) if separator else 22, - username=self.conn.login, password=self.conn.passwd, - timeout=SSH_CONNECT_TIMEOUT_SECONDS, - banner_timeout=SSH_CONNECT_TIMEOUT_SECONDS, - auth_timeout=SSH_CONNECT_TIMEOUT_SECONDS, - channel_timeout=SSH_CONNECT_TIMEOUT_SECONDS, - ) - self.sftp = self.ssh.open_sftp() - self.sftp.get_channel().settimeout(SSH_IO_TIMEOUT_SECONDS) + self._open_sftp() except Exception: self.ssh.close() raise return self - def __exit__(self, *_): + def _connect(self): + hostname, port = self.key[:2] + ssh = paramiko.SSHClient() try: - if self.sftp is not None: - self.sftp.close() + ssh.load_system_host_keys() + ssh.set_missing_host_key_policy(paramiko.RejectPolicy()) + ssh.connect( + hostname=hostname, port=port, + username=self.conn.login, password=self.conn.passwd, + timeout=SSH_CONNECT_TIMEOUT_SECONDS, + banner_timeout=SSH_CONNECT_TIMEOUT_SECONDS, + auth_timeout=SSH_CONNECT_TIMEOUT_SECONDS, + channel_timeout=SSH_CONNECT_TIMEOUT_SECONDS, + ) + except Exception: + ssh.close() + raise + return ssh + + def _open_sftp(self): + # paramiko waits for the SFTP subsystem without a deadline; on a dead + # connection closing it ends the wait with an error + watchdog = threading.Timer(SSH_CONNECT_TIMEOUT_SECONDS, self.ssh.close) + watchdog.daemon = True + watchdog.start() + try: + self.sftp = self.ssh.open_sftp() finally: - if self.ssh is not None: - self.ssh.close() + watchdog.cancel() + transport = self.ssh.get_transport() + if transport is None or not transport.is_active(): + self.sftp.close() + raise paramiko.SSHException(_("The SFTP session did not start in time")) + self.sftp.get_channel().settimeout(SSH_IO_TIMEOUT_SECONDS) + + def __exit__(self, exc_type, *_): + if self.ssh is None: + return + try: + self.sftp.close() + except Exception: + self.ssh.close() + raise + # after an error the connection is not trusted for the next chunk + if exc_type is None: + _ssh_pool.give(self.key, self.ssh, self.opened) + else: + self.ssh.close() def info(self, path): try: