Bind SCP functions

5 files changed, 153 insertions, 8 deletions

diff --git a/libssh2-sys/lib.rs b/libssh2-sys/lib.rs
index f73d5cd..628b397 100644
--- a/libssh2-sys/lib.rs
+++ b/libssh2-sys/lib.rs
@@ -392,4 +392,15 @@ extern {
                                        kind: c_int) -> c_int;
     pub fn libssh2_knownhost_init(sess: *mut LIBSSH2_SESSION)
                                   -> *mut LIBSSH2_KNOWNHOSTS;
+
+    // scp
+    pub fn libssh2_scp_recv(sess: *mut LIBSSH2_SESSION,
+                            path: *const c_char,
+                            sb: *mut libc::stat) -> *mut LIBSSH2_CHANNEL;
+    pub fn libssh2_scp_send64(sess: *mut LIBSSH2_SESSION,
+                              path: *const c_char,
+                              mode: c_int,
+                              size: u64,
+                              mtime: libc::time_t,
+                              atime: libc::time_t) -> *mut LIBSSH2_CHANNEL;
 }
diff --git a/src/channel.rs b/src/channel.rs
index 0d30903..971e72f 100644
--- a/src/channel.rs
+++ b/src/channel.rs
@@ -1,3 +1,4 @@
+use std::cmp;
 use std::io;
 use std::kinds::marker;
 use std::vec;
@@ -9,6 +10,7 @@ pub struct Channel<'a> {
     raw: *mut raw::LIBSSH2_CHANNEL,
     sess: &'a Session,
     marker: marker::NoSync,
+    read_limit: Option<u64>,
 }
 
 /// Data received from when a program exits with a signal.
@@ -53,6 +55,7 @@ impl<'a> Channel<'a> {
             raw: raw,
             sess: sess,
             marker: marker::NoSync,
+            read_limit: None,
         }
     }
 
@@ -86,7 +89,8 @@ impl<'a> Channel<'a> {
 
     /// Check if the remote host has sent an EOF status for the selected stream.
     pub fn eof(&self) -> bool {
-        unsafe { raw::libssh2_channel_eof(self.raw) != 0 }
+        self.read_limit == Some(0) ||
+            unsafe { raw::libssh2_channel_eof(self.raw) != 0 }
     }
 
     /// Initiate a request on a session type channel.
@@ -267,6 +271,15 @@ impl<'a> Channel<'a> {
     /// to be the stderr substream.
     pub fn read_stream(&mut self, stream_id: uint, data: &mut [u8])
                        -> Result<uint, Error> {
+        if self.eof() { return Err(Error::eof()) }
+
+        let data = match self.read_limit {
+            Some(amt) => {
+                let len = data.len();
+                data.slice_to_mut(cmp::min(amt as uint, len))
+            }
+            None => data,
+        };
         unsafe {
             let rc = raw::libssh2_channel_read_ex(self.raw,
                                                   stream_id as c_int,
@@ -274,6 +287,10 @@ impl<'a> Channel<'a> {
                                                   data.len() as size_t);
             if rc < 0 { try!(self.sess.rc(rc)); }
             if rc == 0 && self.eof() { return Err(Error::eof()) }
+            match self.read_limit {
+                Some(ref mut amt) => *amt -= rc as u64,
+                None => {}
+            }
             Ok(rc as uint)
         }
     }
@@ -355,6 +372,12 @@ impl<'a> Channel<'a> {
         try!(self.sess.rc(rc));
         Ok(ret as uint)
     }
+
+    /// Artificially limit the number of bytes that will be read from this
+    /// channel.
+    pub fn limit_read(&mut self, limit: u64) {
+        self.read_limit = Some(limit);
+    }
 }
 
 impl<'a> Writer for Channel<'a> {
diff --git a/src/session.rs b/src/session.rs
index b96fa6e..24f658e 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -1,8 +1,9 @@
+use std::io;
 use std::kinds::marker;
 use std::mem;
 use std::raw as stdraw;
 use std::str;
-use libc::{c_uint, c_int, c_void, c_long};
+use libc::{mod, c_uint, c_int, c_void, c_long};
 
 use {raw, Error, DisconnectCode, ByApplication, SessionFlag, HostKeyType};
 use {MethodType, Agent, Channel, Listener, HashType, KnownHosts};
@@ -482,6 +483,52 @@ impl Session {
         }
     }
 
+    /// Request a file from the remote host via SCP.
+    pub fn scp_recv(&self, path: &Path)
+                    -> Result<(Channel, io::FileStat), Error> {
+        let path = path.to_c_str();
+        unsafe {
+            let mut sb: libc::stat = mem::zeroed();
+            let ret = raw::libssh2_scp_recv(self.raw, path.as_ptr(), &mut sb);
+            if ret.is_null() { return Err(Error::last_error(self).unwrap()) }
+
+            // Hm, apparently when we scp_recv() a file the actual channel
+            // itself does not respond well to read_to_end(), and it also sends
+            // an extra 0 byte (or so it seems). To work around this we
+            // artificially limit the channel to a certain amount of bytes that
+            // can be read.
+            let mut c = Channel::from_raw(self, ret);
+            c.limit_read(sb.st_size as u64);
+            Ok((c, mkstat(&sb)))
+        }
+    }
+
+    /// Send a file to the remote host via SCP.
+    ///
+    /// The `remote_path` provided will the remote file name. The `times`
+    /// argument is a tuple of (mtime, atime), and will default to the remote
+    /// host's current time if not specified.
+    pub fn scp_send(&self, remote_path: &Path, mode: io::FilePermission,
+                    size: u64, times: Option<(u64, u64)>)
+                    -> Result<Channel, Error> {
+        let path = remote_path.to_c_str();
+        let (mtime, atime) = times.unwrap_or((0, 0));
+        unsafe {
+            let ret = raw::libssh2_scp_send64(self.raw,
+                                              path.as_ptr(),
+                                              mode.bits() as c_int,
+                                              size,
+                                              mtime as libc::time_t,
+                                              atime as libc::time_t);
+
+            if ret.is_null() {
+                Err(Error::last_error(self).unwrap())
+            } else {
+                Ok(Channel::from_raw(self, ret))
+            }
+        }
+    }
+
     /// Gain access to the underlying raw libssh2 session pointer.
     pub fn raw(&self) -> *mut raw::LIBSSH2_SESSION { self.raw }
 
@@ -498,6 +545,53 @@ impl Session {
     }
 }
 
+// Sure do wish this was exported in libnative!
+fn mkstat(stat: &libc::stat) -> io::FileStat {
+    #[cfg(windows)] type Mode = libc::c_int;
+    #[cfg(unix)]    type Mode = libc::mode_t;
+
+    // FileStat times are in milliseconds
+    fn mktime(secs: u64, nsecs: u64) -> u64 { secs * 1000 + nsecs / 1000000 }
+
+    #[cfg(not(target_os = "linux"), not(target_os = "android"))]
+    fn flags(stat: &libc::stat) -> u64 { stat.st_flags as u64 }
+    #[cfg(target_os = "linux")] #[cfg(target_os = "android")]
+    fn flags(_stat: &libc::stat) -> u64 { 0 }
+
+    #[cfg(not(target_os = "linux"), not(target_os = "android"))]
+    fn gen(stat: &libc::stat) -> u64 { stat.st_gen as u64 }
+    #[cfg(target_os = "linux")] #[cfg(target_os = "android")]
+    fn gen(_stat: &libc::stat) -> u64 { 0 }
+
+    io::FileStat {
+        size: stat.st_size as u64,
+        kind: match (stat.st_mode as Mode) & libc::S_IFMT {
+            libc::S_IFREG => io::TypeFile,
+            libc::S_IFDIR => io::TypeDirectory,
+            libc::S_IFIFO => io::TypeNamedPipe,
+            libc::S_IFBLK => io::TypeBlockSpecial,
+            libc::S_IFLNK => io::TypeSymlink,
+            _ => io::TypeUnknown,
+        },
+        perm: io::FilePermission::from_bits_truncate(stat.st_mode as u32),
+        created: mktime(stat.st_ctime as u64, stat.st_ctime_nsec as u64),
+        modified: mktime(stat.st_mtime as u64, stat.st_mtime_nsec as u64),
+        accessed: mktime(stat.st_atime as u64, stat.st_atime_nsec as u64),
+        unstable: io::UnstableFileStat {
+            device: stat.st_dev as u64,
+            inode: stat.st_ino as u64,
+            rdev: stat.st_rdev as u64,
+            nlink: stat.st_nlink as u64,
+            uid: stat.st_uid as u64,
+            gid: stat.st_gid as u64,
+            blksize: stat.st_blksize as u64,
+            blocks: stat.st_blocks as u64,
+            flags: flags(stat),
+            gen: gen(stat),
+        }
+    }
+}
+
 impl Drop for Session {
     fn drop(&mut self) {
         unsafe {
diff --git a/tests/channel.rs b/tests/channel.rs
index 2344cce..21a678f 100644
--- a/tests/channel.rs
+++ b/tests/channel.rs
@@ -56,15 +56,9 @@ fn shell() {
 #[test]
 fn setenv() {
     let (_tcp, sess) = ::authed_session();
-    {
     let mut channel = sess.channel_session().unwrap();
     let _ = channel.setenv("FOO", "BAR");
     channel.close().unwrap();
-    drop(channel);
-    }
-    sess.disconnect(None, "lol", None).unwrap();
-    drop(sess);
-    drop(_tcp);
 }
 
 #[test]
diff --git a/tests/session.rs b/tests/session.rs
index 4c8f907..e3275d1 100644
--- a/tests/session.rs
+++ b/tests/session.rs
@@ -1,4 +1,5 @@
 use std::os;
+use std::io::{mod, File, TempDir};
 
 use ssh2::{mod, Session};
 
@@ -48,3 +49,25 @@ fn keepalive() {
     sess.keepalive_set(false, 10).unwrap();
     sess.keepalive_send().unwrap();
 }
+
+#[test]
+fn scp_recv() {
+    let (_tcp, sess) = ::authed_session();
+    let (mut ch, _) = sess.scp_recv(&Path::new(".ssh/authorized_keys")).unwrap();
+    let data = ch.read_to_string().unwrap();
+    let p = Path::new(os::getenv("HOME").unwrap()).join(".ssh/authorized_keys");
+    let expected = File::open(&p).read_to_string().unwrap();
+    assert!(data == expected);
+}
+
+#[test]
+fn scp_send() {
+    let td = TempDir::new("test").unwrap();
+    let (_tcp, sess) = ::authed_session();
+    let mut ch = sess.scp_send(&td.path().join("foo"),
+                               io::UserFile, 6, None).unwrap();
+    ch.write(b"foobar").unwrap();
+    drop(ch);
+    let actual = File::open(&td.path().join("foo")).read_to_end().unwrap();
+    assert_eq!(actual.as_slice(), b"foobar");
+}