From 8bce5882461c905c6565c701bddcf97907fa13e6 Mon Sep 17 00:00:00 2001
From: Compey <hi@devcomp.xyz>
Date: Thu, 20 Jul 2023 17:01:35 +0530
Subject: [PATCH] feat: custom port support for relay

---
 src/cli.rs    |  5 ++++-
 src/main.rs   |  4 ++--
 src/server.rs | 12 +++++++++---
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index cd2c616..370c202 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -29,5 +29,8 @@ pub enum Commands {
         password: String,
     },
     /// Start relay server
-    Relay {},
+    Relay {
+        #[clap(value_parser, required = false)]
+        port: Option<u16>,
+    },
 }
diff --git a/src/main.rs b/src/main.rs
index d79afe9..58d97a3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,8 +28,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
             println!("Receiving password {}", password);
             receive(password).await?
         }
-        Commands::Relay {} => {
-            serve().await?;
+        Commands::Relay { port } => {
+            serve(port.to_owned()).await?;
         }
     }
     Ok(())
diff --git a/src/server.rs b/src/server.rs
index d3f8cf4..d745a96 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -2,6 +2,7 @@ use crate::handshake::Handshake;
 use anyhow::Result;
 use bytes::Bytes;
 use std::collections::HashMap;
+use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio::io::{copy, AsyncWriteExt};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
@@ -87,9 +88,14 @@ impl Client {
     }
 }
 
-pub async fn serve() -> Result<()> {
-    let addr = "127.0.0.1:8080".to_string();
-    let listener = TcpListener::bind(&addr).await?;
+pub async fn serve(port: Option<u16>) -> Result<()> {
+    let port: u16 = match port {
+        Some(port) => port,
+        None => 8080u16,
+    };
+
+    let addr = SocketAddr::from(([127, 0, 0, 1], port));
+    let listener = TcpListener::bind(addr).await?;
     let state = Arc::new(Mutex::new(Shared::new()));
     let (tx, _rx) = broadcast::channel::<Bytes>(100);
     println!("Listening on: {}", addr);