Bind SCP functions

2 files changed, 119 insertions, 2 deletions

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 {