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

src/args.rs

@@ -1,24 +1,29 @@
-use clap::crate_description;
-use clap::{Arg, ArgMatches};
+use clap::{Arg, ArgMatches, Command};
 use rustls::{Certificate, PrivateKey};
-use std::net::{IpAddr, SocketAddr};
+use std::env;
+use std::net::IpAddr;
 use std::path::{Path, PathBuf};
-use std::{env, fs, io};
 
 use crate::auth::parse_auth;
+use crate::tls::{load_certs, load_private_key};
 use crate::BoxResult;
 
-const ABOUT: &str = concat!("\n", crate_description!()); // Add extra newline.
-
-fn app() -> clap::Command<'static> {
-    clap::command!()
-        .about(ABOUT)
+fn app() -> Command<'static> {
+    Command::new(env!("CARGO_CRATE_NAME"))
+        .version(env!("CARGO_PKG_VERSION"))
+        .author(env!("CARGO_PKG_AUTHORS"))
+        .about(concat!(
+            env!("CARGO_PKG_DESCRIPTION"),
+            " - ",
+            env!("CARGO_PKG_REPOSITORY")
+        ))
         .arg(
             Arg::new("address")
                 .short('b')
                 .long("bind")
-                .default_value("0.0.0.0")
                 .help("Specify bind address")
+                .multiple_values(true)
+                .multiple_occurrences(true)
                 .value_name("address"),
         )
         .arg(
@@ -111,7 +116,8 @@ pub fn matches() -> ArgMatches {
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Args {
-    pub addr: SocketAddr,
+    pub addrs: Vec<IpAddr>,
+    pub port: u16,
     pub path: PathBuf,
     pub path_prefix: String,
     pub uri_prefix: String,
@@ -132,9 +138,12 @@ impl Args {
     /// If a parsing error ocurred, exit the process and print out informative
     /// error message to user.
     pub fn parse(matches: ArgMatches) -> BoxResult<Args> {
-        let ip = matches.value_of("address").unwrap_or_default();
         let port = matches.value_of_t::<u16>("port")?;
-        let addr = to_addr(ip, port)?;
+        let addrs = matches
+            .values_of("address")
+            .map(|v| v.collect())
+            .unwrap_or_else(|| vec!["0.0.0.0", "::"]);
+        let addrs: Vec<IpAddr> = Args::parse_addrs(&addrs)?;
         let path = Args::parse_path(matches.value_of_os("path").unwrap_or_default())?;
         let path_prefix = matches
             .value_of("path-prefix")
@@ -166,7 +175,8 @@ impl Args {
         };
 
         Ok(Args {
-            addr,
+            addrs,
+            port,
             path,
             path_prefix,
             uri_prefix,
@@ -182,7 +192,25 @@ impl Args {
         })
     }
 
-    /// Parse path.
+    fn parse_addrs(addrs: &[&str]) -> BoxResult<Vec<IpAddr>> {
+        let mut ip_addrs = vec![];
+        let mut invalid_addrs = vec![];
+        for addr in addrs {
+            match addr.parse::<IpAddr>() {
+                Ok(v) => {
+                    ip_addrs.push(v);
+                }
+                Err(_) => {
+                    invalid_addrs.push(*addr);
+                }
+            }
+        }
+        if !invalid_addrs.is_empty() {
+            return Err(format!("Invalid bind address `{}`", invalid_addrs.join(",")).into());
+        }
+        Ok(ip_addrs)
+    }
+
     fn parse_path<P: AsRef<Path>>(path: P) -> BoxResult<PathBuf> {
         let path = path.as_ref();
         if !path.exists() {
@@ -197,43 +225,3 @@ impl Args {
             .map_err(|err| format!("Failed to access path `{}`: {}", path.display(), err,).into())
     }
 }
-
-fn to_addr(ip: &str, port: u16) -> BoxResult<SocketAddr> {
-    let ip: IpAddr = ip.parse()?;
-    Ok(SocketAddr::new(ip, port))
-}
-
-// Load public certificate from file.
-fn load_certs(filename: &str) -> BoxResult<Vec<Certificate>> {
-    // Open certificate file.
-    let certfile = fs::File::open(&filename)
-        .map_err(|e| format!("Failed to access `{}`, {}", &filename, e))?;
-    let mut reader = io::BufReader::new(certfile);
-
-    // Load and return certificate.
-    let certs = rustls_pemfile::certs(&mut reader).map_err(|_| "Failed to load certificate")?;
-    if certs.is_empty() {
-        return Err("No supported certificate in file".into());
-    }
-    Ok(certs.into_iter().map(Certificate).collect())
-}
-
-// Load private key from file.
-fn load_private_key(filename: &str) -> BoxResult<PrivateKey> {
-    // Open keyfile.
-    let keyfile = fs::File::open(&filename)
-        .map_err(|e| format!("Failed to access `{}`, {}", &filename, e))?;
-    let mut reader = io::BufReader::new(keyfile);
-
-    // Load and return a single private key.
-    let keys = rustls_pemfile::read_all(&mut reader)
-        .map_err(|e| format!("There was a problem with reading private key: {:?}", e))?
-        .into_iter()
-        .find_map(|item| match item {
-            rustls_pemfile::Item::RSAKey(key) | rustls_pemfile::Item::PKCS8Key(key) => Some(key),
-            _ => None,
-        })
-        .ok_or("No supported private key in file")?;
-
-    Ok(PrivateKey(keys))
-}

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);

src/server.rs

@@ -8,7 +8,6 @@ use async_zip::Compression;
 use chrono::{TimeZone, Utc};
 use futures::stream::StreamExt;
 use futures::TryStreamExt;
-use get_if_addrs::get_if_addrs;
 use headers::{
     AcceptRanges, AccessControlAllowHeaders, AccessControlAllowOrigin, ContentLength, ContentRange,
     ContentType, ETag, HeaderMap, HeaderMapExt, IfModifiedSince, IfNoneMatch, IfRange,
@@ -18,28 +17,23 @@ use hyper::header::{
     HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_DISPOSITION, CONTENT_TYPE, ORIGIN, RANGE,
     WWW_AUTHENTICATE,
 };
-use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Method, StatusCode, Uri};
 use percent_encoding::percent_decode;
-use rustls::ServerConfig;
 use serde::Serialize;
-use std::convert::Infallible;
 use std::fs::Metadata;
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::SystemTime;
 use tokio::fs::File;
 use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWrite};
-use tokio::net::TcpListener;
 use tokio::{fs, io};
-use tokio_rustls::TlsAcceptor;
 use tokio_util::codec::{BytesCodec, FramedRead};
 use tokio_util::io::{ReaderStream, StreamReader};
 use uuid::Uuid;
 
-type Request = hyper::Request<Body>;
-type Response = hyper::Response<Body>;
+pub type Request = hyper::Request<Body>;
+pub type Response = hyper::Response<Body>;
 
 const INDEX_HTML: &str = include_str!("../assets/index.html");
 const INDEX_CSS: &str = include_str!("../assets/index.css");
@@ -55,82 +49,34 @@ macro_rules! status {
     };
 }
 
-pub async fn serve(args: Args) -> BoxResult<()> {
-    let args = Arc::new(args);
-    let inner = Arc::new(InnerService::new(args.clone()));
-    match args.tls.clone() {
-        Some((certs, key)) => {
-            let config = ServerConfig::builder()
-                .with_safe_defaults()
-                .with_no_client_auth()
-                .with_single_cert(certs, key)?;
-            let tls_acceptor = TlsAcceptor::from(Arc::new(config));
-            let arc_acceptor = Arc::new(tls_acceptor);
-            let listener = TcpListener::bind(&args.addr).await?;
-            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
-            let incoming =
-                hyper::server::accept::from_stream(incoming.filter_map(|socket| async {
-                    match socket {
-                        Ok(stream) => match arc_acceptor.clone().accept(stream).await {
-                            Ok(val) => Some(Ok::<_, Infallible>(val)),
-                            Err(_) => None,
-                        },
-                        Err(_) => None,
-                    }
-                }));
-            let server = hyper::Server::builder(incoming).serve(make_service_fn(move |_| {
-                let inner = inner.clone();
-                async move {
-                    Ok::<_, Infallible>(service_fn(move |req| {
-                        let inner = inner.clone();
-                        inner.call(req)
-                    }))
-                }
-            }));
-            print_listening(&args.addr, &args.uri_prefix, true);
-            server.await?;
-        }
-        None => {
-            let server = hyper::Server::try_bind(&args.addr)?.serve(make_service_fn(move |_| {
-                let inner = inner.clone();
-                async move {
-                    Ok::<_, Infallible>(service_fn(move |req| {
-                        let inner = inner.clone();
-                        inner.call(req)
-                    }))
-                }
-            }));
-            print_listening(&args.addr, &args.uri_prefix, false);
-            server.await?;
-        }
-    }
-    Ok(())
-}
-
-struct InnerService {
+pub struct Server {
     args: Arc<Args>,
 }
 
-impl InnerService {
+impl Server {
     pub fn new(args: Arc<Args>) -> Self {
         Self { args }
     }
 
-    pub async fn call(self: Arc<Self>, req: Request) -> Result<Response, hyper::Error> {
+    pub async fn call(
+        self: Arc<Self>,
+        req: Request,
+        addr: SocketAddr,
+    ) -> Result<Response, hyper::Error> {
         let method = req.method().clone();
         let uri = req.uri().clone();
         let cors = self.args.cors;
 
         let mut res = match self.handle(req).await {
             Ok(res) => {
-                info!(r#""{} {}" - {}"#, method, uri, res.status());
+                info!(r#"{} "{} {}" - {}"#, addr, method, uri, res.status());
                 res
             }
             Err(err) => {
                 let mut res = Response::default();
                 let status = StatusCode::INTERNAL_SERVER_ERROR;
                 status!(res, status);
-                error!(r#""{} {}" - {} {}"#, method, uri, status, err);
+                error!(r#"{} "{} {}" - {} {}"#, addr, method, uri, status, err);
                 res
             }
         };
@@ -1052,48 +998,6 @@ fn to_content_range(range: &Range, complete_length: u64) -> Option<ContentRange>
     })
 }
 
-fn print_listening(addr: &SocketAddr, prefix: &str, tls: bool) {
-    let prefix = encode_uri(prefix.trim_end_matches('/'));
-    let addrs = retrieve_listening_addrs(addr);
-    let protocol = if tls { "https" } else { "http" };
-    if addrs.len() == 1 {
-        println!("Listening on {}://{}{}", protocol, addr, prefix);
-    } else {
-        let message = addrs
-            .iter()
-            .map(|addr| format!("  {}://{}{}", protocol, addr, prefix))
-            .collect::<Vec<String>>()
-            .join("\n");
-        println!("Listening on:\n{}\n", message);
-    }
-}
-
-fn retrieve_listening_addrs(addr: &SocketAddr) -> Vec<SocketAddr> {
-    let ip = addr.ip();
-    let port = addr.port();
-    if ip.is_unspecified() {
-        if let Ok(interfaces) = get_if_addrs() {
-            let mut ifaces: Vec<IpAddr> = interfaces
-                .into_iter()
-                .map(|v| v.ip())
-                .filter(|v| {
-                    if ip.is_ipv4() {
-                        v.is_ipv4()
-                    } else {
-                        v.is_ipv6()
-                    }
-                })
-                .collect();
-            ifaces.sort();
-            return ifaces
-                .into_iter()
-                .map(|v| SocketAddr::new(v, port))
-                .collect();
-        }
-    }
-    vec![addr.to_owned()]
-}
-
 fn encode_uri(v: &str) -> String {
     let parts: Vec<_> = v.split('/').map(urlencoding::encode).collect();
     parts.join("/")

src/tls.rs

@@ -0,0 +1,158 @@
+use core::task::{Context, Poll};
+use futures::ready;
+use hyper::server::accept::Accept;
+use hyper::server::conn::{AddrIncoming, AddrStream};
+use rustls::{Certificate, PrivateKey};
+use std::future::Future;
+use std::net::SocketAddr;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::{fs, io};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio_rustls::rustls::ServerConfig;
+
+enum State {
+    Handshaking(tokio_rustls::Accept<AddrStream>),
+    Streaming(tokio_rustls::server::TlsStream<AddrStream>),
+}
+
+// tokio_rustls::server::TlsStream doesn't expose constructor methods,
+// so we have to TlsAcceptor::accept and handshake to have access to it
+// TlsStream implements AsyncRead/AsyncWrite handshaking tokio_rustls::Accept first
+pub struct TlsStream {
+    state: State,
+    remote_addr: SocketAddr,
+}
+
+impl TlsStream {
+    fn new(stream: AddrStream, config: Arc<ServerConfig>) -> TlsStream {
+        let remote_addr = stream.remote_addr();
+        let accept = tokio_rustls::TlsAcceptor::from(config).accept(stream);
+        TlsStream {
+            state: State::Handshaking(accept),
+            remote_addr,
+        }
+    }
+    pub fn remote_addr(&self) -> SocketAddr {
+        self.remote_addr
+    }
+}
+
+impl AsyncRead for TlsStream {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context,
+        buf: &mut ReadBuf,
+    ) -> Poll<io::Result<()>> {
+        let pin = self.get_mut();
+        match pin.state {
+            State::Handshaking(ref mut accept) => match ready!(Pin::new(accept).poll(cx)) {
+                Ok(mut stream) => {
+                    let result = Pin::new(&mut stream).poll_read(cx, buf);
+                    pin.state = State::Streaming(stream);
+                    result
+                }
+                Err(err) => Poll::Ready(Err(err)),
+            },
+            State::Streaming(ref mut stream) => Pin::new(stream).poll_read(cx, buf),
+        }
+    }
+}
+
+impl AsyncWrite for TlsStream {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        let pin = self.get_mut();
+        match pin.state {
+            State::Handshaking(ref mut accept) => match ready!(Pin::new(accept).poll(cx)) {
+                Ok(mut stream) => {
+                    let result = Pin::new(&mut stream).poll_write(cx, buf);
+                    pin.state = State::Streaming(stream);
+                    result
+                }
+                Err(err) => Poll::Ready(Err(err)),
+            },
+            State::Streaming(ref mut stream) => Pin::new(stream).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.state {
+            State::Handshaking(_) => Poll::Ready(Ok(())),
+            State::Streaming(ref mut stream) => Pin::new(stream).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match self.state {
+            State::Handshaking(_) => Poll::Ready(Ok(())),
+            State::Streaming(ref mut stream) => Pin::new(stream).poll_shutdown(cx),
+        }
+    }
+}
+
+pub struct TlsAcceptor {
+    config: Arc<ServerConfig>,
+    incoming: AddrIncoming,
+}
+
+impl TlsAcceptor {
+    pub fn new(config: Arc<ServerConfig>, incoming: AddrIncoming) -> TlsAcceptor {
+        TlsAcceptor { config, incoming }
+    }
+}
+
+impl Accept for TlsAcceptor {
+    type Conn = TlsStream;
+    type Error = io::Error;
+
+    fn poll_accept(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Conn, Self::Error>>> {
+        let pin = self.get_mut();
+        match ready!(Pin::new(&mut pin.incoming).poll_accept(cx)) {
+            Some(Ok(sock)) => Poll::Ready(Some(Ok(TlsStream::new(sock, pin.config.clone())))),
+            Some(Err(e)) => Poll::Ready(Some(Err(e))),
+            None => Poll::Ready(None),
+        }
+    }
+}
+
+// Load public certificate from file.
+pub fn load_certs(filename: &str) -> Result<Vec<Certificate>, Box<dyn std::error::Error>> {
+    // Open certificate file.
+    let certfile = fs::File::open(&filename)
+        .map_err(|e| format!("Failed to access `{}`, {}", &filename, e))?;
+    let mut reader = io::BufReader::new(certfile);
+
+    // Load and return certificate.
+    let certs = rustls_pemfile::certs(&mut reader).map_err(|_| "Failed to load certificate")?;
+    if certs.is_empty() {
+        return Err("No supported certificate in file".into());
+    }
+    Ok(certs.into_iter().map(Certificate).collect())
+}
+
+// Load private key from file.
+pub fn load_private_key(filename: &str) -> Result<PrivateKey, Box<dyn std::error::Error>> {
+    // Open keyfile.
+    let keyfile = fs::File::open(&filename)
+        .map_err(|e| format!("Failed to access `{}`, {}", &filename, e))?;
+    let mut reader = io::BufReader::new(keyfile);
+
+    // Load and return a single private key.
+    let keys = rustls_pemfile::read_all(&mut reader)
+        .map_err(|e| format!("There was a problem with reading private key: {:?}", e))?
+        .into_iter()
+        .find_map(|item| match item {
+            rustls_pemfile::Item::RSAKey(key) | rustls_pemfile::Item::PKCS8Key(key) => Some(key),
+            _ => None,
+        })
+        .ok_or("No supported private key in file")?;
+
+    Ok(PrivateKey(keys))
+}