resolve merge conflicts

src/main.rs

@@ -12,6 +12,7 @@ use {
         borrow::Cow,
         error::Error,
         ffi::OsStr,
+        fmt::Write,
         fs::File,
         io::BufReader,
         net::SocketAddr,
@@ -45,8 +46,14 @@ fn main() -> Result {
             let (stream, _) = listener.accept().await?;
             let arc = mimetypes.clone();
             tokio::spawn(async {
-                if let Err(e) = handle_request(stream, arc).await {
+                match RequestHandle::new(stream, arc).await {
-                    log::error!("{:?}", e);
+                    Ok(handle) => match handle.handle().await {
+                        Ok(info) => log::info!("{}", info),
+                        Err(err) => log::warn!("{}", err),
+                    },
+                    Err(log_line) => {
+                        log::warn!("{}", log_line);
+                    }
                 }
             });
         }
@@ -71,6 +78,7 @@ struct Args {
     language: Option<String>,
     silent: bool,
     serve_secret: bool,
+    log_ips: bool,
 }
 
 fn args() -> Result<Args> {
@@ -115,6 +123,7 @@ fn args() -> Result<Args> {
     opts.optflag("s", "silent", "Disable logging output");
     opts.optflag("h", "help", "Print this help menu");
     opts.optflag("", "serve-secret", "Enable serving secret files (files/directories starting with a dot)");
+    opts.optflag("", "log-ip", "Output IP addresses when logging");
 
     let matches = opts.parse(&args[1..]).map_err(|f| f.to_string())?;
     if matches.opt_present("h") {
@@ -144,6 +153,7 @@ fn args() -> Result<Args> {
         language: matches.opt_str("lang"),
         silent: matches.opt_present("s"),
         serve_secret: matches.opt_present("serve-secret"),
+    	log_ips: matches.opt_present("log-ip"),
     })
 }
 
@@ -155,18 +165,6 @@ fn check_path(s: String) -> Result<String, String> {
     }
 }
 
-/// Handle a single client session (request + response).
-async fn handle_request(stream: TcpStream, mimetypes: Arc<RwLock<FileOptions>>) -> Result {
-    let stream = &mut TLS.accept(stream).await?;
-
-    match parse_request(stream).await {
-        Ok(url) => send_response(url, stream, mimetypes).await?,
-        Err((status, msg)) => send_header(stream, status, &[msg]).await?,
-    }
-    stream.shutdown().await?;
-    Ok(())
-}
-
 /// TLS configuration.
 static TLS: Lazy<TlsAcceptor> = Lazy::new(|| acceptor().unwrap());
 
@@ -182,10 +180,57 @@ fn acceptor() -> Result<TlsAcceptor> {
     Ok(TlsAcceptor::from(Arc::new(config)))
 }
 
+struct RequestHandle {
+    stream: TlsStream<TcpStream>,
+    log_line: String,
+    metadata: Arc<RwLock<FileOptions>>,
+}
+
+impl RequestHandle {
+    /// Creates a new request handle for the given stream. If establishing the TLS
+    /// session fails, returns a corresponding log line.
+    async fn new(stream: TcpStream, metadata: Arc<RwLock<FileOptions>>) -> Result<Self, String> {
+        let log_line = format!(
+            "{} {}",
+            stream.local_addr().unwrap(),
+            if ARGS.log_ips {
+                stream
+                    .peer_addr()
+                    .expect("could not get peer address")
+                    .to_string()
+            } else {
+                // Do not log IP address, but something else so columns still line up.
+                "-".into()
+            }
+        );
+
+        match TLS.accept(stream).await {
+            Ok(stream) => Ok(Self { stream, log_line, metadata }),
+            Err(e) => Err(format!("{} error:{}", log_line, e)),
+        }
+    }
+
+    /// Do the necessary actions to handle this request. Returns a corresponding
+    /// log line as Err or Ok, depending on if the request finished with or
+    /// without errors.
+    async fn handle(mut self) -> Result<String, String> {
+        // not already in error condition
+        let result = match self.parse_request().await {
+            Ok(url) => self.send_response(url).await,
+            Err((status, msg)) => self.send_header(status, msg).await,
+        };
+
+        if let Err(e) = result {
+            Err(format!("{} error:{}", self.log_line, e))
+        } else if let Err(e) = self.stream.shutdown().await {
+            Err(format!("{} error:{}", self.log_line, e))
+        } else {
+            Ok(self.log_line)
+        }
+    }
+
     /// Return the URL requested by the client.
-async fn parse_request(
+    async fn parse_request(&mut self) -> std::result::Result<Url, (u8, &'static str)> {
-    stream: &mut TlsStream<TcpStream>,
-) -> std::result::Result<Url, (u8, &'static str)> {
         // Because requests are limited to 1024 bytes (plus 2 bytes for CRLF), we
         // can use a fixed-sized buffer on the stack, avoiding allocations and
         // copying, and stopping bad clients from making us use too much memory.
@@ -195,7 +240,8 @@ async fn parse_request(
 
         // Read until CRLF, end-of-stream, or there's no buffer space left.
         loop {
-        let bytes_read = stream
+            let bytes_read = self
+                .stream
                 .read(buf)
                 .await
                 .or(Err((59, "Request ended unexpectedly")))?;
@@ -207,8 +253,11 @@ async fn parse_request(
             }
             buf = &mut request[len..];
         }
-    let request = std::str::from_utf8(&request[..len - 2]).or(Err((59, "Non-UTF-8 request")))?;
+        let request =
-    log::info!("Got request for {:?}", request);
+            std::str::from_utf8(&request[..len - 2]).or(Err((59, "Non-UTF-8 request")))?;
+
+        // log literal request (might be different from or not an actual URL)
+        write!(self.log_line, " \"{}\"", request).unwrap();
 
         let url = Url::parse(request).or(Err((59, "Invalid URL")))?;
 
@@ -217,14 +266,16 @@ async fn parse_request(
             return Err((53, "Unsupported URL scheme"));
         }
         // TODO: Can be simplified by https://github.com/servo/rust-url/pull/651
-    if let (Some(Host::Domain(expected)), Some(Host::Domain(host))) = (url.host(), &ARGS.hostname) {
+        if let (Some(Host::Domain(expected)), Some(Host::Domain(host))) =
+            (url.host(), &ARGS.hostname)
+        {
             if host != expected {
                 return Err((53, "Proxy request refused"));
             }
         }
         if let Some(port) = url.port() {
             // Validate that the port in the URL is the same as for the stream this request came in on.
-        if port != stream.get_ref().0.local_addr().unwrap().port() {
+            if port != self.stream.get_ref().0.local_addr().unwrap().port() {
                 return Err((53, "proxy request refused"));
             }
         }
@@ -232,17 +283,14 @@ async fn parse_request(
     }
 
     /// Send the client the file located at the requested URL.
-async fn send_response(
+    async fn send_response(&mut self, url: Url) -> Result {
-    url: Url,
-    stream: &mut TlsStream<TcpStream>,
-    mimetypes: Arc<RwLock<FileOptions>>,
-) -> Result {
         let mut path = std::path::PathBuf::from(&ARGS.content_dir);
         if let Some(segments) = url.path_segments() {
             for segment in segments {
                 if !ARGS.serve_secret && segment.starts_with('.') {
                     // Do not serve anything that looks like a hidden file.
-                return send_header(stream, 52, &["If I told you, it would not be a secret."])
+                    return self
+                        .send_header(52, "If I told you, it would not be a secret.")
                         .await;
                 }
                 path.push(&*percent_decode_str(segment).decode_utf8()?);
@@ -257,13 +305,13 @@ async fn send_response(
                     path.push("index.gmi");
                     if !path.exists() && path.with_file_name(".directory-listing-ok").exists() {
                         path.pop();
-                    return list_directory(stream, &path).await;
+                        return self.list_directory(&path).await;
                     }
                 } else {
                     // if client is not redirected, links may not work as expected without trailing slash
                     let mut url = url;
                     url.set_path(&format!("{}/", url.path()));
-                return send_header(stream, 31, &[url.as_str()]).await;
+                    return self.send_header(31, url.as_str()).await;
                 }
             }
         }
@@ -272,41 +320,49 @@ async fn send_response(
         let mut file = match tokio::fs::File::open(&path).await {
             Ok(file) => file,
             Err(e) => {
-            send_header(stream, 51, &["Not found, sorry."]).await?;
+                self.send_header(51, "Not found, sorry.").await?;
                 Err(e)?
             }
         };
 
         // Send header.
-    let mut locked = mimetypes.write().await;
+        let mut locked = self.metadata.write().await;
         let data = locked.get(&path);
-    if data.is_empty() || data.starts_with(";") {
+        let mime = if data.is_empty() || data.starts_with(';') {
             // guess MIME type
             if path.extension() == Some(OsStr::new("gmi")) {
-            send_header(stream, 20, &["text/gemini", data]).await?;
+            	format!("text/gemini{}", data)
             } else {
                 let mime = mime_guess::from_path(&path).first_or_octet_stream();
-            send_header(stream, 20, &[mime.essence_str(), data]).await?;
+                format!("{}{}", mime.essence_str(), data)
-        };
+            }
         } else {
             // this must be a full MIME type
-        send_header(stream, 20, &[data]).await?;
+        	data.to_owned()
-    }
+        };
         drop(locked);
+        self.send_header(20, &mime).await?;
 
         // Send body.
-    tokio::io::copy(&mut file, stream).await?;
+        tokio::io::copy(&mut file, &mut self.stream).await?;
         Ok(())
     }
 
-async fn list_directory(stream: &mut TlsStream<TcpStream>, path: &Path) -> Result {
+    async fn list_directory(&mut self, path: &Path) -> Result {
         // https://url.spec.whatwg.org/#path-percent-encode-set
-    const ENCODE_SET: AsciiSet = CONTROLS.add(b' ')
+        const ENCODE_SET: AsciiSet = CONTROLS
-        .add(b'"').add(b'#').add(b'<').add(b'>')
+            .add(b' ')
-        .add(b'?').add(b'`').add(b'{').add(b'}');
+            .add(b'"')
+            .add(b'#')
+            .add(b'<')
+            .add(b'>')
+            .add(b'?')
+            .add(b'`')
+            .add(b'{')
+            .add(b'}');
 
         log::info!("Listing directory {:?}", path);
-    send_header(stream, 20, &["text/gemini"]).await?;
+        self.send_header(20, "text/gemini").await?;
         let mut entries = tokio::fs::read_dir(path).await?;
         let mut lines = vec![];
         while let Some(entry) = entries.next_entry().await? {
@@ -328,18 +384,18 @@ async fn list_directory(stream: &mut TlsStream<TcpStream>, path: &Path) -> Resul
         }
         lines.sort();
         for line in lines {
-        stream.write_all(line.as_bytes()).await?;
+            self.stream.write_all(line.as_bytes()).await?;
         }
         Ok(())
     }
 
-async fn send_header(stream: &mut TlsStream<TcpStream>, status: u8, meta: &[&str]) -> Result {
+    async fn send_header(&mut self, status: u8, meta: &str) -> Result {
-    use std::fmt::Write;
+        // add response status and response meta
-    let mut response = String::with_capacity(64);
+        write!(self.log_line, " {} \"{}\"", status, meta)?;
-    write!(response, "{} ", status)?;
+
-    response.extend(meta.iter().copied());
+        self.stream
-    log::info!("Responding with status {:?}", response);
+            .write_all(format!("{} {}\r\n", status, meta).as_bytes())
-    response.push_str("\r\n");
+            .await?;
-    stream.write_all(response.as_bytes()).await?;
         Ok(())
     }
+}

src/metadata.rs

@@ -84,7 +84,7 @@ impl FileOptions {
                 // discard any I/O errors
                 .filter_map(|line| line.ok())
                 // filter out comment lines
-                .filter(|line| !line.trim_start().starts_with("#"))
+                .filter(|line| !line.trim_start().starts_with('#'))
                 .for_each(|line| {
                     // split line at colon
                     let parts = line.splitn(2, ':').collect::<Vec<_>>();