Merge pull request #756 from justin-stephenson/retry_ssh_again

Retry on libssh `SSH_AGAIN` return code

Author: webknjaz

docs/changelog-fragments/756.feature.1.rst

@@ -0,0 +1,8 @@
+Added a ``pylibsshext.session.connect()`` parameter
+``open_session_retries`` -- by :user:`justin-stephenson`.
+
+The ``open_session_retries`` session ``connect()``
+parameter allows a configurable number of retries if
+libssh ``ssh_channel_open_session()`` returns ``SSH_AGAIN``.
+The default option value is 0, no retries will be
+attempted.

docs/changelog-fragments/756.feature.2.rst

@@ -0,0 +1,8 @@
+Added a ``pylibsshext.session.connect()`` parameter
+``timeout_usec`` to set ``SSH_OPTIONS_TIMEOUT_USEC``.
+
+This allows setting the ``SSH_OPTIONS_TIMEOUT_USEC``
+ssh option, though ``SSH_OPTIONS_TIMEOUT`` is a more
+practical option.
+
+-- by :user:`justin-stephenson`

src/pylibsshext/channel.pxd

@@ -23,6 +23,7 @@ cdef class Channel:
     cdef  _session
     cdef libssh.ssh_channel _libssh_channel
     cdef libssh.ssh_session _libssh_session
+    cdef _open_session_with_retries(self, libssh.ssh_channel test_channel)
 
 cdef class ChannelCallback:
     cdef callbacks.ssh_channel_callbacks_struct callback

src/pylibsshext/channel.pyx

@@ -24,7 +24,7 @@ from libc.string cimport memset
 
 from pylibsshext.errors cimport LibsshChannelException
 from pylibsshext.errors import LibsshChannelReadFailure
-from pylibsshext.session cimport get_libssh_session
+from pylibsshext.session cimport get_libssh_session, get_session_retries
 
 from subprocess import CompletedProcess
 
@@ -63,12 +63,8 @@ cdef class Channel:
 
         if self._libssh_channel is NULL:
             raise MemoryError
-        rc = libssh.ssh_channel_open_session(self._libssh_channel)
 
-        if rc != libssh.SSH_OK:
-            libssh.ssh_channel_free(self._libssh_channel)
-            self._libssh_channel = NULL
-            raise LibsshChannelException("Failed to open_session: [%d]" % rc)
+        self._open_session_with_retries(self._libssh_channel)
 
     def __dealloc__(self):
         if self._libssh_channel is not NULL:
@@ -158,16 +154,28 @@ cdef class Channel:
             response = recv_buff.getvalue()
         return response
 
+    cdef _open_session_with_retries(self, libssh.ssh_channel channel):
+        retry = get_session_retries(self._session)
+
+        for attempt in range(retry + 1):
+            rc = libssh.ssh_channel_open_session(channel)
+            if rc == libssh.SSH_OK:
+                break
+            if rc == libssh.SSH_AGAIN and attempt < retry:
+                continue
+            # either SSH_ERROR, or SSH_AGAIN with final attempt
+            if rc != libssh.SSH_OK:
+                self._libssh_channel = NULL
+                libssh.ssh_channel_free(channel)
+                raise LibsshChannelException(f"Failed to open_session: [{rc}]")
+
     def exec_command(self, command):
         # request_exec requires a fresh channel each run, so do not use the existing channel
         cdef libssh.ssh_channel channel = libssh.ssh_channel_new(self._libssh_session)
         if channel is NULL:
             raise MemoryError
 
-        rc = libssh.ssh_channel_open_session(channel)
-        if rc != libssh.SSH_OK:
-            libssh.ssh_channel_free(channel)
-            raise LibsshChannelException("Failed to open_session: [{0}]".format(rc))
+        self._open_session_with_retries(channel)
 
         result = CompletedProcess(args=command, returncode=-1, stdout=b'', stderr=b'')
 

src/pylibsshext/session.pxd

@@ -26,6 +26,8 @@ cdef class Session:
     cdef _hash_py
     cdef _fingerprint_py
     cdef _keytype_py
+    cdef _retries
     cdef _channel_callbacks
 
 cdef libssh.ssh_session get_libssh_session(Session session)
+cdef int get_session_retries(Session session)

src/pylibsshext/session.pyx

@@ -31,6 +31,7 @@ OPTS_MAP = {
     "user": libssh.SSH_OPTIONS_USER,
     "port": libssh.SSH_OPTIONS_PORT,
     "timeout": libssh.SSH_OPTIONS_TIMEOUT,
+    "timeout_usec": libssh.SSH_OPTIONS_TIMEOUT_USEC,
     "knownhosts": libssh.SSH_OPTIONS_KNOWNHOSTS,
     "proxycommand": libssh.SSH_OPTIONS_PROXYCOMMAND,
     "key_exchange_algorithms": libssh.SSH_OPTIONS_KEY_EXCHANGE,
@@ -108,6 +109,7 @@ cdef class Session(object):
         self._hash_py = None
         self._fingerprint_py = None
         self._keytype_py = None
+        self._retries = 0
         # Due to delayed freeing of channels, some older libssh versions might expect
         # the callbacks to be around even after we free the underlying channels so
         # we should free them only when we terminate the session.
@@ -175,7 +177,7 @@ cdef class Session(object):
         elif key == "port":
             value_uint = value
             libssh.ssh_options_set(self._libssh_session, key_m, &value_uint)
-        elif key == "timeout":
+        elif key in {"timeout", "timeout_usec"}:
             value_long = value
             libssh.ssh_options_set(self._libssh_session, key_m, &value_long)
         else:
@@ -235,9 +237,17 @@ cdef class Session(object):
         file should be validated. It defaults to True
         :type host_key_checking: boolean
 
+        :param open_session_retries: The number of retries to attempt when libssh
+                                      channel function ``ssh_channel_open_session()`` returns ``SSH_AGAIN``. It defaults
+                                      to 0, no retries attempted.
+        :type open_session_retries: integer
+
         :param timeout: The timeout in seconds for the TCP connect
         :type timeout: long integer
 
+        :param timeout_usec: The timeout in microseconds for the TCP connect
+        :type timeout_usec: long integer
+
         :param port: The ssh server port to connect to
         :type port: integer
 
@@ -261,6 +271,9 @@ cdef class Session(object):
                 libssh.ssh_disconnect(self._libssh_session)
                 raise
 
+        if 'open_session_retries' in kwargs:
+            self._retries = kwargs['open_session_retries']
+
         # We need to userauth_none before we can query the available auth types
         rc = libssh.ssh_userauth_none(self._libssh_session, NULL)
         if rc == libssh.SSH_AUTH_SUCCESS:
@@ -553,3 +566,6 @@ cdef class Session(object):
 
 cdef libssh.ssh_session get_libssh_session(Session session):
     return session._libssh_session
+
+cdef int get_session_retries(Session session):
+    return session._retries

tests/_service_utils.py

@@ -69,6 +69,7 @@ def wait_for_svc_ready_state(
 
 def ensure_ssh_session_connected(  # noqa: WPS317
         ssh_session, sshd_addr, ssh_clientkey_path,  # noqa: WPS318
+        ssh_session_retries=0,
 ):
     """Attempt connecting to the SSH server until successful.
 
@@ -89,4 +90,5 @@ def ensure_ssh_session_connected(  # noqa: WPS317
         private_key=ssh_clientkey_path.read_bytes(),
         host_key_checking=False,
         look_for_keys=False,
+        open_session_retries=ssh_session_retries,
     )

tests/conftest.py

@@ -126,6 +126,25 @@ def ssh_client_session(ssh_session_connect):
         del ssh_session  # noqa: WPS420
 
 
+@pytest.fixture
+def ssh_session_connect_retries(sshd_addr, ssh_clientkey_path):
+    """
+    Authenticate existing session object against SSHD with a private SSH key.
+
+    This sets ``open_session_retries=100`` and it returns a function
+    that takes session as parameter.
+
+    :returns: Function that will connect the session.
+    :rtype: Callback
+    """
+    return partial(
+        ensure_ssh_session_connected,
+        sshd_addr=sshd_addr,
+        ssh_clientkey_path=ssh_clientkey_path,
+        ssh_session_retries=100,
+    )
+
+
 @pytest.fixture
 def ssh_session_connect(sshd_addr, ssh_clientkey_path):
     """

tests/unit/channel_test.py

@@ -8,12 +8,17 @@
 
 import pytest
 
+from pylibsshext.errors import LibsshChannelException
 from pylibsshext.session import Session
 
 
 COMMAND_TIMEOUT = 30
 POLL_EXIT_CODE_TIMEOUT = 5
 POLL_TIMEOUT = 5000
+# Lowest possible timeout_usec is 1000 because timeout_usec / 1000 must be > 0,
+# otherwise libssh sets timeout to 10000
+SMALL_TIMEOUT_USEC = 1000
+LARGE_TIMEOUT_SEC = 10
 
 
 @pytest.fixture
@@ -33,6 +38,42 @@ def ssh_channel(ssh_client_session):
         chan.close()
 
 
+def test_open_session_small_timeout(ssh_session_connect):
+    """Test opening a new channel with a small timeout value.
+
+    This generates an exception from ``ssh_channel_open_session()``
+    returning ``SSH_AGAIN`` with the ``usec`` timeout and default
+    ``open_session_retries`` value of ``0``.
+    """
+    ssh_session = Session()
+    ssh_session_connect(ssh_session)
+    ssh_session.set_ssh_options('timeout_usec', SMALL_TIMEOUT_USEC)
+    error_msg = '^Failed to open_session'
+    with pytest.raises(LibsshChannelException, match=error_msg):
+        ssh_session.new_channel()
+    ssh_session.close()
+
+
+def test_open_session_large_timeout(ssh_session_connect):
+    """Test opening a new channel with a large timeout value."""
+    ssh_session = Session()
+    ssh_session_connect(ssh_session)
+    ssh_session.set_ssh_options('timeout', LARGE_TIMEOUT_SEC)
+    ssh_channel = ssh_session.new_channel()
+    ssh_channel.close()
+    ssh_session.close()
+
+
+def test_open_session_small_timeout_with_retries(ssh_session_connect_retries):
+    """Test with a small timeout value and retries set."""
+    ssh_session = Session()
+    ssh_session_connect_retries(ssh_session)
+    ssh_session.set_ssh_options('timeout_usec', SMALL_TIMEOUT_USEC)
+    ssh_channel = ssh_session.new_channel()
+    ssh_channel.close()
+    ssh_session.close()
+
+
 def exec_second_command(ssh_channel):
     """Check the standard output of ``exec_command()`` as a string."""
     u_cmd = ssh_channel.exec_command('echo -n Hello Again')