From 65c379e501f05f8aca4cb4db52d93b80ad45f7cd Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Sun, 17 Dec 2023 19:26:58 +0800
Subject: [PATCH] bind argument

---
 src/main.rs   | 11 +++++++++++
 src/server.rs | 12 +++++++++---
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index d8c1122..75ad85d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -54,6 +54,17 @@ fn main() {
                 .multiple(true)
                 .default_value("")
         )
+        .arg(
+            Arg::with_name("bind")
+                .help("bind to the interface associated with the address <host>")
+                .takes_value(true)
+                .long("bind")
+                .short("B")
+                .value_name("host")
+                .required(true)
+                .multiple(true)
+                .default_value("0.0.0.0")
+        )
         .arg(
             Arg::with_name("debug")
                 .help("emit debug-level logging on stderr; default is info and above")
diff --git a/src/server.rs b/src/server.rs
index dd2ae56..ccba606 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -20,7 +20,7 @@
 
 use std::error::Error;
 use std::io;
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
+use std::net::{IpAddr, Ipv4Addr, Shutdown, SocketAddr};
 use std::sync::atomic::{AtomicBool, AtomicU16, Ordering};
 use std::sync::mpsc::channel;
 use std::sync::{Arc, Mutex};
@@ -401,13 +401,19 @@ pub fn serve(args: ArgMatches) -> BoxResult<()> {
         log::debug!("limiting service to {} concurrent clients", client_limit);
     }
 
+    let addr = args
+        .value_of("bind")
+        .unwrap_or("0.0.0.0")
+        .parse::<IpAddr>()
+        .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED));
+
     //start listening for connections
     let port: u16 = args.value_of("port").unwrap().parse()?;
     let listener: TcpListener = if args.is_present("version6") {
-        TcpListener::bind(&SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port))
+        TcpListener::bind(&SocketAddr::new(addr, port))
             .unwrap_or_else(|_| panic!("failed to bind TCP socket, port {}", port))
     } else {
-        TcpListener::bind(&SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port))
+        TcpListener::bind(&SocketAddr::new(addr, port))
             .unwrap_or_else(|_| panic!("failed to bind TCP socket, port {}", port))
     };
     log::info!("server listening on {}", listener.local_addr()?);