client can set remote address and server address if desired

README.md

@@ -27,17 +27,17 @@ client <-> server <-> remote
 
 	Example:
 	```
-	host = '192.168.1.1'
-	port = 6000
-	remote = '127.0.0.1:1080'
+	host = '192.168.1.1'         # host of server
+	port = 8022                  # port of server
+	remote = '127.0.0.1:1080'    # default remote address (can be customized per client)
 	```
 
 2. Generate server keypair by running `portguard gen-key -c config.toml`.
 
 	After that, `config.toml` becomes:
 	```
 	host = '192.168.1.1'
-	port = 6000
+	port = 8022
 	remote = '127.0.0.1:1080'
 	pubkey = '1y3HW8TDxChtke5nyEdLGj+OkQSg8JjLdalSHzD+aWI='
 	prikey = 'eHg7jR/IZwEZEqeyR27IUTN0py5a3+wP0uM+z9HeWn8='
@@ -48,7 +48,7 @@ client <-> server <-> remote
 	After that, `config.toml` becomes:
 	```
 	host = '192.168.1.1'
-	port = 6000
+	port = 8022
 	remote = '127.0.0.1:1080'
 	pubkey = '1y3HW8TDxChtke5nyEdLGj+OkQSg8JjLdalSHzD+aWI='
 	prikey = 'eHg7jR/IZwEZEqeyR27IUTN0py5a3+wP0uM+z9HeWn8='
@@ -61,15 +61,15 @@ client <-> server <-> remote
 
 3. Run `portguard server -c config.toml` on server
 
-4. Run `./pgcli`  on client without any configs
-(local port can be customized with `./pgcli -p port` if you like).
+4. Run `pgcli` on client without any configs
+(local port or server address can be customized with `./pgcli -p port -s saddr:sport` if you like).
 
 5. All TCP traffic to client's local port is forwarded to remote by server with encryption.
 
 ## TODO
 
 - [ ] I'm not familar with Noise protocol, now in my code every connection between client and server needs to handshake.
-- [ ] Set remote address per client.
+- [x] Set remote address per client.
 - [ ] Plan to use other Noise implementation.
 - [ ] Improve performance
 

src/client.rs

@@ -1,5 +1,5 @@
 use crate::consts::{CONF_BUF_LEN, PATTERN};
-use crate::proxy::transfer;
+use crate::proxy;
 use futures::FutureExt;
 use log;
 use serde::{Deserialize, Serialize};
@@ -30,15 +30,22 @@ impl Client {
     pub fn new(port: u16) -> Client {
         Client { port }
     }
-    pub async fn run_client_proxy(self) -> Result<(), Box<dyn Error>> {
+    pub async fn run_client_proxy(self, server: Option<SocketAddr>) -> Result<(), Box<dyn Error>> {
         let this = Arc::new(self);
-        let conf: ClientConfig = bincode::deserialize(&CLIENT_CONF_BUF)?;
+        let mut conf: ClientConfig = bincode::deserialize(&CLIENT_CONF_BUF)?;
+        // overwrite server address
+        if let Some(server_addr) = server {
+            conf.server_addr = server_addr;
+        }
         let shared_conf = Arc::new(conf);
         let listen_addr: SocketAddr = format!("127.0.0.1:{}", this.port).parse()?;
         log::info!("Client listening on: {:?}", listen_addr);
         log::info!("Portguard server on: {:?}", shared_conf.server_addr);
         log::info!("Target address: {:?}", shared_conf.target_addr);
-        log::debug!("Portguard server public key: {:?}", base64::encode(&shared_conf.server_pubkey));
+        log::debug!(
+            "Portguard server public key: {:?}",
+            base64::encode(&shared_conf.server_pubkey)
+        );
 
         let listener = TcpListener::bind(listen_addr).await?;
 
@@ -54,7 +61,11 @@ impl Client {
         Ok(())
     }
 
-    async fn handle_connection(&self, inbound: TcpStream, conf: &ClientConfig) -> Result<(), Box<dyn Error>> {
+    async fn handle_connection(
+        &self,
+        inbound: TcpStream,
+        conf: &ClientConfig,
+    ) -> Result<(), Box<dyn Error>> {
         log::info!("New incoming peer_addr {:?}", inbound.peer_addr());
         let initiator = snowstorm::Builder::new(PATTERN.parse()?)
             .remote_public_key(&conf.server_pubkey)
@@ -63,7 +74,7 @@ impl Client {
         let outbound = TcpStream::connect(conf.server_addr).await?;
         let enc_outbound = NoiseStream::handshake(outbound, initiator).await?;
 
-        let transfer = transfer(inbound, enc_outbound).map(|r| {
+        let transfer = proxy::transfer(inbound, enc_outbound).map(|r| {
             if let Err(e) = r {
                 log::error!("Transfer error occured. error={}", e);
             }

src/main.rs

@@ -23,6 +23,9 @@ struct ClientArgs {
     /// local port to listen
     #[clap(short, long, default_value_t = 6000)]
     port: u16,
+    /// replace another server address in this run
+    #[clap(short, long)]
+    server: Option<String>
 }
 
 #[derive(Subcommand)]
@@ -31,21 +34,28 @@ enum Commands {
     Client(ClientArgs),
     /// run server
     Server {
-        /// config file
+        /// location of config file
         #[clap(short, long)]
         config: PathBuf,
     },
     /// generate client binary
     GenCli {
+        /// location of config file
         #[clap(short, long)]
         config: PathBuf,
+        /// location of output binary
         #[clap(short, long)]
         output: PathBuf,
+        /// name of client
         #[clap(short, long, default_value = "user")]
         name: String,
+        /// client specified remote address
+        #[clap(short, long)]
+        remote: Option<String>,
     },
     /// generate keypairs
     GenKey {
+        /// location of config file
         #[clap(short, long)]
         config: PathBuf,
     },
@@ -61,9 +71,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let cli = Cli::parse();
     let client_cmd = cli.command.unwrap_or(Commands::Client(cli.client));
     match client_cmd {
-        Commands::Client(ClientArgs { port }) => {
+        Commands::Client(ClientArgs { port, server }) => {
             let client = Client::new(port);
-            client.run_client_proxy().await?;
+            let server = server.and_then(|s| s.parse().ok());
+            client.run_client_proxy(server).await?;
         }
         Commands::Server { config: path } => {
             let content = std::fs::read_to_string(&path)?;
@@ -75,12 +86,13 @@ async fn main() -> Result<(), Box<dyn Error>> {
             config: path,
             output: out_path,
             name,
+            remote
         } => {
             let content = std::fs::read_to_string(&path)?;
             let config = toml::de::from_str(&content)?;
 
             let mut server = portguard::server::Server::new(config, &path);
-            server.gen_client(out_path, name)?;
+            server.gen_client(out_path, name, remote)?;
         }
         Commands::GenKey { config: path } => {
             let content = std::fs::read_to_string(&path)?;

src/server.rs

@@ -1,6 +1,6 @@
 use crate::client::{ClientConfig, CLIENT_CONF_BUF};
 use crate::consts::{CONF_BUF_LEN, PATTERN};
-use crate::proxy::transfer;
+use crate::proxy;
 
 use futures::FutureExt;
 use log;
@@ -40,6 +40,8 @@ struct ClientEntry {
     /// client public key for auth
     #[serde(with = "base64_serde")]
     pubkey: Vec<u8>,
+    /// client specified remote address
+    remote: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -50,7 +52,7 @@ pub struct ServerConfig {
     /// server listen port
     #[serde(default = "default_port")]
     pub port: u16,
-    /// remote address hope to proxy
+    /// default remote address hope to proxy
     #[serde(default = "default_remote")]
     remote: String,
     /// server public key
@@ -59,21 +61,21 @@ pub struct ServerConfig {
     /// server private key
     #[serde(with = "base64_serde", default)]
     prikey: Vec<u8>,
-    #[serde(default)]
+    #[serde(serialize_with = "toml::ser::tables_last", default)]
     /// infomation of clients
     clients: HashMap<String, ClientEntry>,
 }
 
 fn default_port() -> u16 {
-    6000
+    8022
 }
 
 fn default_host() -> String {
     "192.168.1.1".to_string()
 }
 
 fn default_remote() -> String {
-    "127.0.0.1:8080".to_string()
+    "127.0.0.1:1080".to_string()
 }
 
 impl ServerConfig {
@@ -97,16 +99,24 @@ impl Server {
         }
     }
     // fn from_config()
-    pub fn gen_client<P: AsRef<Path>>(&mut self, exe_path: P, username: String) -> Result<(), Box<dyn Error>> {
+    pub fn gen_client<P: AsRef<Path>>(
+        &mut self,
+        exe_path: P,
+        username: String,
+        remote: Option<String>,
+    ) -> Result<(), Box<dyn Error>> {
         let mut cli_conf: ClientConfig = bincode::deserialize(&CLIENT_CONF_BUF)?;
         log::debug!("Previous static variable: {:?}", cli_conf);
         // 1. set client config
         let key = snowstorm::Builder::new(PATTERN.parse()?).generate_keypair()?;
         cli_conf.client_prikey = key.private;
-        log::debug!("Client private key: {:?}", base64::encode(&cli_conf.client_prikey));
+        log::debug!(
+            "Client private key: {:?}",
+            base64::encode(&cli_conf.client_prikey)
+        );
         cli_conf.server_pubkey = self.config.pubkey.clone();
         cli_conf.server_addr = format!("{}:{}", self.config.host, self.config.port).parse()?;
-        cli_conf.target_addr = self.config.remote.parse()?;
+        cli_conf.target_addr = remote.as_ref().unwrap_or(&self.config.remote).parse()?;
         log::debug!("New client static variable: {:?}", cli_conf);
         // 2. crate new binary
         let exe = env::current_exe()?;
@@ -130,7 +140,11 @@ impl Server {
             fs::remove_file(&new_exe)?;
         }
         // 4. add new client to server config
-        let client = ClientEntry { name: username, pubkey: key.public };
+        let client = ClientEntry {
+            name: username,
+            pubkey: key.public,
+            remote,
+        };
         let ent = self.config.clients.entry(base64::encode(&client.pubkey));
         ent.or_insert(client);
         // 5. save server config
@@ -151,27 +165,22 @@ impl Server {
     pub async fn run_server_proxy(self) -> Result<(), Box<dyn Error>> {
         let this = Arc::new(self);
         let listen_addr: SocketAddr = format!("0.0.0.0:{}", this.config.port).parse().unwrap();
-        let out_addr: SocketAddr = this.config.remote.parse().unwrap();
         log::info!("Listening on port: {:?}", listen_addr);
 
         let listener = TcpListener::bind(listen_addr).await?;
 
         while let Ok((inbound, _)) = listener.accept().await {
             let this = Arc::clone(&this);
             tokio::spawn(async move {
-                if let Err(e) = this.handle_connection(inbound, out_addr).await {
+                if let Err(e) = this.handle_connection(inbound).await {
                     log::warn!("{}", e);
                 }
             });
         }
         Ok(())
     }
 
-    async fn handle_connection(
-        &self,
-        inbound: TcpStream,
-        out_addr: SocketAddr,
-    ) -> Result<(), Box<dyn Error>> {
+    async fn handle_connection(&self, inbound: TcpStream) -> Result<(), Box<dyn Error>> {
         log::info!("New incoming peer_addr {:?}", inbound.peer_addr());
         let responder = snowstorm::Builder::new(PATTERN.parse()?)
             .local_private_key(&self.config.prikey)
@@ -181,8 +190,16 @@ impl Server {
             self.config.clients.contains_key(&token)
         })
         .await?;
+        // at this point, client already passed verification
+        // if it specifies a remote address, use it
+        // can use `.unwrap()` here because client must have a static key
+        let token = base64::encode(enc_inbound.get_state().get_remote_static().unwrap());
+        let remote = self.config.clients.get(&token).unwrap().remote.as_ref();
+        let out_addr: SocketAddr = remote.unwrap_or(&self.config.remote).parse()?;
+
+        log::info!("Start proxying {:?} to {:?}", enc_inbound.get_inner().peer_addr(), out_addr);
         let outbound = TcpStream::connect(out_addr).await?;
-        let transfer = transfer(enc_inbound, outbound).map(|r| {
+        let transfer = proxy::transfer(enc_inbound, outbound).map(|r| {
             if let Err(e) = r {
                 log::error!("Transfer error occured. error={}", e);
             }
@@ -195,9 +212,6 @@ impl Server {
 fn serialize_conf_to_buf(conf: &ClientConfig) -> Result<[u8; CONF_BUF_LEN], Box<dyn Error>> {
     let v = &bincode::serialize(&conf)?;
     let mut bytes: [u8; CONF_BUF_LEN] = [0; CONF_BUF_LEN];
-    // for i in 0..v.len() {
-        // bytes[i] = v[i];
-    // }
     bytes[..v.len()].clone_from_slice(&v[..]);
     Ok(bytes)
 }