From f784e8f87818da0718c704c0457804b1f5eb5e22 Mon Sep 17 00:00:00 2001
From: Ganesh B Nalawade <ganesh634@gmail.com>
Date: Sat, 11 Jul 2020 10:02:51 +0530
Subject: [PATCH] Fix exec_command issue

Fixes https://github.com/ansible/pylibssh/issues/57
Fixes https://github.com/ansible/pylibssh/issues/56

*  Use channel created with Channel class
   instead of creating a new channel with
   exec_command()
---
 src/pylibsshext/channel.pyx | 28 +++++++++++-----------------
 tests/unit/channel_test.py  |  6 ------
 2 files changed, 11 insertions(+), 23 deletions(-)

diff --git a/src/pylibsshext/channel.pyx b/src/pylibsshext/channel.pyx
index 5062e82dc..93d7ca99c 100644
--- a/src/pylibsshext/channel.pyx
+++ b/src/pylibsshext/channel.pyx
@@ -61,8 +61,9 @@ cdef class Channel:
             raise LibsshChannelException("Failed to open_session: [%d]" % rc)
 
     def __dealloc__(self):
-        libssh.ssh_channel_free(self._libssh_channel)
-        self._libssh_channel = NULL
+        if self._libssh_channel is not NULL:
+            libssh.ssh_channel_free(self._libssh_channel)
+            self._libssh_channel = NULL
 
     def request_shell(self):
         self.request_pty()
@@ -133,17 +134,10 @@ cdef class Channel:
         return response
 
     def exec_command(self, command):
-        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 CalledProcessError()
+        rc = libssh.ssh_channel_request_exec(self._libssh_channel, command.encode("utf-8"))
 
-        rc = libssh.ssh_channel_request_exec(channel, command.encode("utf-8"))
         if rc != libssh.SSH_OK:
-            libssh.ssh_channel_free(channel)
+            libssh.ssh_channel_free(self._libssh_channel)
             raise CalledProcessError()
 
         cdef callbacks.ssh_channel_callbacks_struct cb
@@ -152,12 +146,11 @@ cdef class Channel:
         result = CompletedProcess(args=command, returncode=-1, stdout=b'', stderr=b'')
         cb.userdata = <void *>result
         callbacks.ssh_callbacks_init(&cb)
-        callbacks.ssh_set_channel_callbacks(channel, &cb)
+        callbacks.ssh_set_channel_callbacks(self._libssh_channel, &cb)
 
-        libssh.ssh_channel_send_eof(channel)
-        result.returncode = libssh.ssh_channel_get_exit_status(channel)
+        libssh.ssh_channel_send_eof(self._libssh_channel)
 
-        libssh.ssh_channel_free(channel)
+        result.returncode = self.get_channel_exit_status()
 
         return result
 
@@ -165,5 +158,6 @@ cdef class Channel:
         return libssh.ssh_channel_get_exit_status(self._libssh_channel)
 
     def close(self):
-        libssh.ssh_channel_free(self._libssh_channel)
-        self._libssh_channel = NULL
+        if self._libssh_channel is not NULL:
+            libssh.ssh_channel_free(self._libssh_channel)
+            self._libssh_channel = NULL
diff --git a/tests/unit/channel_test.py b/tests/unit/channel_test.py
index 30e649c3c..3a27b45b6 100644
--- a/tests/unit/channel_test.py
+++ b/tests/unit/channel_test.py
@@ -26,12 +26,6 @@ def ssh_channel(ssh_client_session):
         chan.close()
 
 
-@pytest.mark.xfail(
-    reason='This test causes SEGFAULT, flakily. '
-    'Ref: https://github.com/ansible/pylibssh/issues/57',  # noqa: WPS326
-    run=False,
-    strict=False,
-)
 @pytest.mark.forked
 def test_exec_command(ssh_channel):
     """Test getting the output of a remotely executed command."""