Skip to content

Commit

Permalink
Better handling for exit in pty shells
Browse files Browse the repository at this point in the history
Handled a condition where pty shells would hang on exit. This is
currently resolved by explicitly passing "exit" to the channel when
exiting the context manager. I'm sure there are better ways to handle
this, but this works for now.
  • Loading branch information
JacobCallahan committed Apr 4, 2024
1 parent 3c77122 commit ea2d305
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 2 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "hussh"
version = "0.1.3"
version = "0.1.4"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ dev = [
"pexpect",
"pre-commit",
"pytest",
"pytest-randomly",
"ruff",
]

Expand Down
8 changes: 7 additions & 1 deletion src/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,7 @@ impl Connection {
channel.shell().unwrap();
Ok(InteractiveShell {
channel: ChannelWrapper { channel },
pty: pty.unwrap_or(false),
exit_result: None,
})
}
Expand All @@ -518,16 +519,18 @@ pub struct ChannelWrapper {
#[derive(Clone)]
struct InteractiveShell {
channel: ChannelWrapper,
pty: bool,
#[pyo3(get)]
exit_result: Option<SSHResult>,
}

#[pymethods]
impl InteractiveShell {
#[new]
fn new(channel: ChannelWrapper) -> Self {
fn new(channel: ChannelWrapper, pty: bool) -> Self {
InteractiveShell {
channel,
pty,
exit_result: None,
}
}
Expand Down Expand Up @@ -569,6 +572,9 @@ impl InteractiveShell {
_exc_value: Option<&Bound<'_, PyAny>>,
_traceback: Option<&Bound<'_, PyAny>>,
) -> PyResult<()> {
if self.pty {
self.send("exit\n".to_string(), Some(false)).unwrap();
}
self.exit_result = Some(self.read());
Ok(())
}
Expand Down
10 changes: 10 additions & 0 deletions tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,16 @@ def test_shell_context(conn):
assert sh.exit_result.status != 0


def test_pty_shell_context(conn):
"""Test that we can run multiple commands in a pty shell context."""
with conn.shell(pty=True) as sh:
sh.send("echo test shell")
sh.send("bad command")
assert "test shell" in sh.exit_result.stdout
assert "command not found" in sh.exit_result.stdout
assert sh.exit_result.status != 0


def test_connection_timeout():
"""Test that we can trigger a timeout on connect."""
with pytest.raises(TimeoutError):
Expand Down

0 comments on commit ea2d305

Please sign in to comment.