feat: listen both ipv4 and ipv6 by default (#40)

src/main.rs

@@ -1,17 +1,28 @@
 mod args;
 mod auth;
 mod server;
+mod tls;
 
 #[macro_use]
 extern crate log;
 
-pub type BoxResult<T> = Result<T, Box<dyn std::error::Error>>;
-
-use std::env;
-use std::io::Write;
-
 use crate::args::{matches, Args};
-use crate::server::serve;
+use crate::server::{Request, Server};
+use crate::tls::{TlsAcceptor, TlsStream};
+
+use std::io::Write;
+use std::net::{IpAddr, SocketAddr, TcpListener as StdTcpListener};
+use std::{env, sync::Arc};
+
+use futures::future::join_all;
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
+
+use hyper::server::conn::{AddrIncoming, AddrStream};
+use hyper::service::{make_service_fn, service_fn};
+use rustls::ServerConfig;
+
+pub type BoxResult<T> = Result<T, Box<dyn std::error::Error>>;
 
 #[tokio::main]
 async fn main() {
@@ -30,9 +41,18 @@ async fn run() -> BoxResult<()> {
         .init();
 
     let args = Args::parse(matches())?;
+    let args = Arc::new(args);
+    let handles = serve(args.clone())?;
+    print_listening(args)?;
+
     tokio::select! {
-        ret = serve(args) => {
-            ret
+        ret = join_all(handles) => {
+            for r in ret {
+                if let Err(e) = r {
+                    error!("{}", e);
+                }
+            }
+            Ok(())
         },
         _ = shutdown_signal() => {
             Ok(())
@@ -40,6 +60,121 @@ async fn run() -> BoxResult<()> {
     }
 }
 
+fn serve(args: Arc<Args>) -> BoxResult<Vec<JoinHandle<Result<(), hyper::Error>>>> {
+    let inner = Arc::new(Server::new(args.clone()));
+    let mut handles = vec![];
+    let port = args.port;
+    for ip in args.addrs.iter() {
+        let inner = inner.clone();
+        let incoming = create_addr_incoming(SocketAddr::new(*ip, port))
+            .map_err(|e| format!("Failed to bind `{}:{}`, {}", ip, port, e))?;
+        let serv_func = move |remote_addr: SocketAddr| {
+            let inner = inner.clone();
+            async move {
+                Ok::<_, hyper::Error>(service_fn(move |req: Request| {
+                    let inner = inner.clone();
+                    inner.call(req, remote_addr)
+                }))
+            }
+        };
+        match args.tls.clone() {
+            Some((certs, key)) => {
+                let config = ServerConfig::builder()
+                    .with_safe_defaults()
+                    .with_no_client_auth()
+                    .with_single_cert(certs, key)?;
+                let config = Arc::new(config);
+                let accepter = TlsAcceptor::new(config.clone(), incoming);
+                let new_service = make_service_fn(move |socket: &TlsStream| {
+                    let remote_addr = socket.remote_addr();
+                    serv_func(remote_addr)
+                });
+                let server = tokio::spawn(hyper::Server::builder(accepter).serve(new_service));
+                handles.push(server);
+            }
+            None => {
+                let new_service = make_service_fn(move |socket: &AddrStream| {
+                    let remote_addr = socket.remote_addr();
+                    serv_func(remote_addr)
+                });
+                let server = tokio::spawn(hyper::Server::builder(incoming).serve(new_service));
+                handles.push(server);
+            }
+        };
+    }
+    Ok(handles)
+}
+
+fn create_addr_incoming(addr: SocketAddr) -> BoxResult<AddrIncoming> {
+    use socket2::{Domain, Protocol, Socket, Type};
+    let socket = Socket::new(Domain::for_address(addr), Type::STREAM, Some(Protocol::TCP))?;
+    if addr.is_ipv6() {
+        socket.set_only_v6(true)?;
+    }
+    socket.set_reuse_address(true)?;
+    socket.bind(&addr.into())?;
+    socket.listen(1024 /* Default backlog */)?;
+    let std_listener = StdTcpListener::from(socket);
+    std_listener.set_nonblocking(true)?;
+    let incoming = AddrIncoming::from_listener(TcpListener::from_std(std_listener)?)?;
+    Ok(incoming)
+}
+
+fn print_listening(args: Arc<Args>) -> BoxResult<()> {
+    let mut addrs = vec![];
+    let (mut ipv4, mut ipv6) = (false, false);
+    for ip in args.addrs.iter() {
+        if ip.is_unspecified() {
+            if ip.is_ipv6() {
+                ipv6 = true;
+            } else {
+                ipv4 = true;
+            }
+        } else {
+            addrs.push(*ip);
+        }
+    }
+    if ipv4 || ipv6 {
+        let ifaces = get_if_addrs::get_if_addrs()
+            .map_err(|e| format!("Failed to get local interface addresses: {}", e))?;
+        for iface in ifaces.into_iter() {
+            let local_ip = iface.ip();
+            if ipv4 && local_ip.is_ipv4() {
+                addrs.push(local_ip)
+            }
+            if ipv6 && local_ip.is_ipv6() {
+                addrs.push(local_ip)
+            }
+        }
+    }
+    addrs.sort_unstable();
+    let urls = addrs
+        .into_iter()
+        .map(|addr| match addr {
+            IpAddr::V4(_) => format!("{}:{}", addr, args.port),
+            IpAddr::V6(_) => format!("[{}]:{}", addr, args.port),
+        })
+        .map(|addr| match &args.tls {
+            Some(_) => format!("https://{}", addr),
+            None => format!("http://{}", addr),
+        })
+        .map(|url| format!("{}{}", url, args.uri_prefix))
+        .collect::<Vec<_>>();
+
+    if urls.len() == 1 {
+        println!("Listening on {}", urls[0]);
+    } else {
+        let info = urls
+            .iter()
+            .map(|v| format!("  {}", v))
+            .collect::<Vec<String>>()
+            .join("\n");
+        println!("Listening on:\n{}\n", info);
+    }
+
+    Ok(())
+}
+
 fn handle_err<T>(err: Box<dyn std::error::Error>) -> T {
     eprintln!("error: {}", err);
     std::process::exit(1);