session: Add 'open_session_retries' option

Improve pylibssh handling when libssh ssh_channel_open_session()
returns SSH_AGAIN. Add a new 'open_session_retries' session connect()
parameter to allow a configurable number of retries. SSH_AGAIN may be
returned when setting a low SSH options timeout value.

The default option value is 0, no retries will be attempted.

Author: justin-stephenson

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,20 @@ 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)
+        retry = get_session_retries(session)
+
+        for attempt in range(retry + 1):
+            rc = libssh.ssh_channel_open_session(self._libssh_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:
+                libssh.ssh_channel_free(self._libssh_channel)
+                self._libssh_channel = NULL
+                raise LibsshChannelException("Failed to open_session: [%d]" % rc)
 
     def __dealloc__(self):
         if self._libssh_channel is not NULL:
@@ -164,10 +172,18 @@ cdef class Channel:
         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))
+        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:
+                libssh.ssh_channel_free(channel)
+                raise LibsshChannelException("Failed to open_session: [{0}]".format(rc))
 
         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

@@ -109,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.
@@ -236,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
 
@@ -262,6 +271,9 @@ cdef class Session(object):
                 libssh.ssh_disconnect(self._libssh_session)
                 raise
 
+        if kwargs.get('open_session_retries'):
+            self._retries = kwargs.get('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:
@@ -554,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.
 
@@ -88,5 +89,5 @@ def ensure_ssh_session_connected(  # noqa: WPS317
         user=getpass.getuser(),
         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,24 @@ 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
+    with ssh_session_retries parameter set to '10'
+
+    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=10,
+    )
+
 @pytest.fixture
 def ssh_session_connect(sshd_addr, ssh_clientkey_path):
     """

tests/unit/channel_test.py

@@ -8,6 +8,7 @@
 
 import pytest
 
+from pylibsshext.errors import LibsshChannelException
 from pylibsshext.session import Session
 
 
@@ -32,6 +33,28 @@ def ssh_channel(ssh_client_session):
     finally:
         chan.close()
 
+def test_open_session_timeout(ssh_session_connect):
+    """Test that setting a low timeout value generates an
+       exception from ssh_channel_open_session() with default
+       open_session_retries value of 0.
+    """
+    ssh_session = Session()
+    ssh_session_connect(ssh_session)
+    ssh_session.set_ssh_options("timeout_usec", 10000)
+    error_msg = '^Failed to open_session'
+    with pytest.raises(LibsshChannelException, match=error_msg):
+        ssh_channel = ssh_session.new_channel()
+
+
+def test_open_session_with_retries(ssh_session_connect_retries):
+    """Test with a low timeout value and 'open_session_retries=10'
+        set, ssh_channel_open_session() will succeed.
+    """
+    ssh_session = Session()
+    ssh_session_connect_retries(ssh_session)
+    ssh_session.set_ssh_options("timeout_usec", 10000)
+    ssh_channel = ssh_session.new_channel()
+
 
 def exec_second_command(ssh_channel):
     """Check the standard output of ``exec_command()`` as a string."""