export block directions

4 files changed, 44 insertions, 3 deletions

diff --git a/libssh2-sys/lib.rs b/libssh2-sys/lib.rs
index 3c73afb..ed10da5 100644
--- a/libssh2-sys/lib.rs
+++ b/libssh2-sys/lib.rs
@@ -185,6 +185,9 @@ pub const LIBSSH2_CHANNEL_EXTENDED_DATA_NORMAL: c_int = 0;
 pub const LIBSSH2_CHANNEL_EXTENDED_DATA_IGNORE: c_int = 1;
 pub const LIBSSH2_CHANNEL_EXTENDED_DATA_MERGE: c_int = 2;
 
+pub const LIBSSH2_SESSION_BLOCK_INBOUND: c_int = 1;
+pub const LIBSSH2_SESSION_BLOCK_OUTBOUND: c_int = 2;
+
 pub enum LIBSSH2_SESSION {}
 pub enum LIBSSH2_AGENT {}
 pub enum LIBSSH2_CHANNEL {}
@@ -362,6 +365,7 @@ extern "C" {
     );
     pub fn libssh2_keepalive_send(sess: *mut LIBSSH2_SESSION, seconds_to_next: *mut c_int)
         -> c_int;
+    pub fn libssh2_session_block_directions(sess: *mut LIBSSH2_SESSION) -> c_int;
 
     // agent
     pub fn libssh2_agent_init(sess: *mut LIBSSH2_SESSION) -> *mut LIBSSH2_AGENT;
diff --git a/src/lib.rs b/src/lib.rs
index 4c57256..8f04390 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -147,7 +147,7 @@ pub use error::Error;
 pub use knownhosts::{Host, Hosts, KnownHosts};
 pub use listener::Listener;
 use session::SessionInner;
-pub use session::{KeyboardInteractivePrompt, Prompt, ScpFileStat, Session};
+pub use session::{KeyboardInteractivePrompt, Prompt, ScpFileStat, Session, BlockDirections};
 pub use sftp::{File, FileStat, FileType, OpenType};
 pub use sftp::{OpenFlags, RenameFlags, Sftp};
 pub use DisconnectCode::{AuthCancelledByUser, TooManyConnections};
diff --git a/src/session.rs b/src/session.rs
index d1b3da2..364ddd7 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -92,6 +92,19 @@ pub struct ScpFileStat {
     stat: libc::stat,
 }
 
+/// The io direction an application has to wait for in order not to block.
+#[derive(Debug, PartialEq)]
+pub enum BlockDirections {
+    /// No direction blocked.
+    None,
+    /// Inbound direction blocked.
+    Inbound,
+    /// Outbound direction blockd.
+    Outbound,
+    /// Inbound and Outbound direction blocked.
+    Both,
+}
+
 impl Session {
     /// Initializes an SSH session object.
     ///
@@ -914,6 +927,21 @@ impl Session {
     pub fn rc(&self, rc: c_int) -> Result<(), Error> {
         self.inner.rc(rc)
     }
+
+    /// Returns the blocked io directions that the application needs to wait for.
+    ///
+    /// This function should be used after an error of type `WouldBlock` is returned to
+    /// find out the socket events the application has to wait for.
+    pub fn block_directions(&self) -> BlockDirections {
+        let dir = unsafe { raw::libssh2_session_block_directions(self.inner.raw) };
+        match dir {
+            raw::LIBSSH2_SESSION_BLOCK_INBOUND => BlockDirections::Inbound,
+            raw::LIBSSH2_SESSION_BLOCK_OUTBOUND => BlockDirections::Outbound,
+            x if x == raw::LIBSSH2_SESSION_BLOCK_INBOUND | raw::LIBSSH2_SESSION_BLOCK_OUTBOUND
+              => BlockDirections::Both,
+            _ => BlockDirections::None,
+        }
+    }
 }
 
 impl SessionInner {
diff --git a/tests/all/session.rs b/tests/all/session.rs
index 484610f..86587b4 100644
--- a/tests/all/session.rs
+++ b/tests/all/session.rs
@@ -1,10 +1,10 @@
 use std::env;
 use std::fs::File;
-use std::io::prelude::*;
+use std::io::{self, prelude::*};
 use std::path::Path;
 use tempdir::TempDir;
 
-use ssh2::{HashType, KeyboardInteractivePrompt, MethodType, Prompt, Session};
+use ssh2::{HashType, KeyboardInteractivePrompt, MethodType, Prompt, Session, BlockDirections};
 
 #[test]
 fn session_is_send() {
@@ -177,3 +177,12 @@ fn scp_send() {
         .unwrap();
     assert_eq!(actual, b"foobar");
 }
+
+#[test]
+fn block_directions() {
+    let mut sess = ::authed_session();
+    sess.set_blocking(false);
+    let actual = sess.handshake().map_err(|e| io::Error::from(e).kind());
+    assert_eq!(actual, Err(io::ErrorKind::WouldBlock));
+    assert_eq!(sess.block_directions(), BlockDirections::Inbound);
+}