server: proper parsing of SSH_COMMAND

2 files changed, 30 insertions, 11 deletions

diff --git a/src/common.rs b/src/common.rs
index 175c4b8..bf3910e 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -9,16 +9,6 @@ use std::io::{Read, Write, BufRead, BufReader};
 use std::process::exit;
 use std::net;
 
-// This function is an ugly hack fall-back. It tries to find public IP address of this machine by
-// creating an outdoing connection
-pub fn get_local_ip() -> io::Result<net::IpAddr> {
-    
-    let dummy = try!(net::TcpListener::bind("0.0.0.0:0"));
-    let addr = try!(dummy.local_addr()).ip();
-    drop(dummy);
-    Ok(addr)
-}
-
 pub fn source_files<S: Read + Write>(stream: &mut S, file_path: &str, recursive: bool) {
     println!("SOURCE FILE: {}", file_path);
     if recursive { unimplemented!(); }
diff --git a/src/server.rs b/src/server.rs
index 66dfc69..feb5d09 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -6,6 +6,8 @@ use super::common;
 use std::str::{self, FromStr};
 use std::env;
 use std::net;
+use std::io;
+use std::error::Error;
 use std::env::home_dir;
 use std::process::exit;
 use getopts::Options;
@@ -14,13 +16,40 @@ use crypto::{SecretStream, key2string, string2key, nonce2string, string2nonce};
 use udt_extras::{UdtStream};
 use sodiumoxide::crypto::secretbox;
 
+pub fn get_local_ip() -> Result<net::IpAddr, String> {
+    let ip_str = match env::var("SSH_CONNECTION") {
+        Ok(val) => {
+            match val.split(' ').nth(2) {
+                Some(x) => x.to_string(),
+                None => { return Err(format!("Failed to parse $SSH_CONNECTION: {}", val.to_string())); },
+            }
+        },
+        Err(_) => {
+            println!("Can't find $SSH_CONNECTION; running locally? Falling back to 127.0.0.1");
+            "127.0.0.1".to_string()
+        },
+    };
+
+    // First try IPv4
+    match net::Ipv4Addr::from_str(&ip_str) {
+        Ok(x) => { return Ok(net::IpAddr::V4(x)) },
+        Err(_) => (),
+    };
+    // Then IPv6
+    match net::Ipv6Addr::from_str(&ip_str) {
+        Ok(x) => { return Ok(net::IpAddr::V6(x)) },
+        Err(e) => { return Err(e.description().to_string()); },
+    };
+}
+
+
 fn run_server(path: &str, is_recv: bool, recursive: bool, daemonize: bool, no_crypto: bool) {
 
     // TODO: try to detect the address the SSH connection came in on via the SSH_CONNECTION env
     // variable.
 
     // Connect to an hypothetical local server running on port 61000
-    let listen_addr = common::get_local_ip().unwrap();
+    let listen_addr = get_local_ip().unwrap();
     let port = 61000;
 
     // Accept connection from anybody