move tcpstream assignment to its own function

The recent move to take ownership of TcpStream exposed an issue with
the `handshake` method: if the stream is non-blocking then it may
take several attempts to handshake, but only the first one is able
to transfer ownership.

My initial thought was just to make the TcpStream a required parameter
to `new`, but we have some tests that work with known hosts and the
ssh agent that don't require a tcpstream.

I'm going to review those and see if there is a cleaner overall
solution, but that will likely require more substantial API changes.

For now, the simplest change is to add a separate `set_tcp_stream`
function to make the stream ownership transfer explicit and distinct
from the handshake.

Refs: https://github.com/alexcrichton/ssh2-rs/issues/17

2 files changed, 6 insertions, 3 deletions

diff --git a/tests/all/main.rs b/tests/all/main.rs
index 75f2a29..a021b4c 100644
--- a/tests/all/main.rs
+++ b/tests/all/main.rs
@@ -23,7 +23,8 @@ pub fn authed_session() -> ssh2::Session {
     let user = env::var("USER").unwrap();
     let socket = socket();
     let mut sess = ssh2::Session::new().unwrap();
-    sess.handshake(socket).unwrap();
+    sess.set_tcp_stream(socket);
+    sess.handshake().unwrap();
     assert!(!sess.authenticated());
 
     {
diff --git a/tests/all/session.rs b/tests/all/session.rs
index 5e45d53..cf64f1d 100644
--- a/tests/all/session.rs
+++ b/tests/all/session.rs
@@ -30,7 +30,8 @@ fn smoke_handshake() {
     let user = env::var("USER").unwrap();
     let socket = ::socket();
     let mut sess = Session::new().unwrap();
-    sess.handshake(socket).unwrap();
+    sess.set_tcp_stream(socket);
+    sess.handshake().unwrap();
     sess.host_key().unwrap();
     let methods = sess.auth_methods(&user).unwrap();
     assert!(methods.contains("publickey"), "{}", methods);
@@ -52,7 +53,8 @@ fn keyboard_interactive() {
     let user = env::var("USER").unwrap();
     let socket = ::socket();
     let mut sess = Session::new().unwrap();
-    sess.handshake(socket).unwrap();
+    sess.set_tcp_stream(socket);
+    sess.handshake().unwrap();
     sess.host_key().unwrap();
     let methods = sess.auth_methods(&user).unwrap();
     assert!(methods.contains("keyboard-interactive"), "{}", methods);