Added ability to customize certificate path

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Added
 - `document` API for creating Gemini documents
+- preliminary timeout API by [@Alch-Emi](https://github.com/Alch-Emi)
+- `Response::success_with_body` by [@Alch-Emi](https://github.com/Alch-Emi)
+- `redirect_temporary_lossy` for `Response` and `ResponseHeader`
+- `bad_request_lossy` for `Response` and `ResponseHeader`
+- support for a lot more mime-types in `guess_mime_from_path`, backed by the `mime_guess` crate
 
 ## [0.3.0] - 2020-11-14
 ### Added
@@ -25,4 +30,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.2.0] - 2020-11-14
 ### Added
-- Access to client certificates by [@Alch-Emi](https://github.com/Alch-Emi)
+- Access to client certificates by [@Alch-Emi](https://github.com/Alch-Emi)

Cargo.toml

@@ -21,6 +21,7 @@ itertools = "0.9.0"
 log = "0.4.11"
 webpki = "0.21.0"
 lazy_static = "1.4.0"
+mime_guess = "2.0.3"
 
 [dev-dependencies]
 env_logger = "0.8.1"

examples/certificates.rs

@@ -37,8 +37,10 @@ fn handle_request(users: Arc<RwLock<HashMap<CertBytes, String>>>, request: Reque
             if let Some(user) = users_read.get(cert_bytes) {
                 // The user has already registered
                 Ok(
-                    Response::success(&GEMINI_MIME)
+                    Response::success_with_body(
-                        .with_body(format!("Welcome {}!", user))
+                        &GEMINI_MIME,
+                        format!("Welcome {}!", user)
+                    )
                 )
             } else {
                 // The user still needs to register
@@ -49,11 +51,13 @@ fn handle_request(users: Arc<RwLock<HashMap<CertBytes, String>>>, request: Reque
                     let mut users_write = users.write().await;
                     users_write.insert(cert_bytes.clone(), username.to_owned());
                     Ok(
-                        Response::success(&GEMINI_MIME)
+                        Response::success_with_body(
-                            .with_body(format!(
+                            &GEMINI_MIME,
+                            format!(
                                 "Your account has been created {}!  Welcome!",
                                 username
-                            ))
+                            )
+                        )
                     )
                 } else {
                     // The user didn't provide input, and should be prompted

src/lib.rs

@@ -6,12 +6,14 @@ use std::{
     io::BufReader,
     sync::Arc,
     path::PathBuf,
+    time::Duration,
 };
 use futures::{future::BoxFuture, FutureExt};
 use tokio::{
     prelude::*,
     io::{self, BufStream},
     net::{TcpStream, ToSocketAddrs},
+    time::timeout,
 };
 use tokio::net::TcpListener;
 use rustls::ClientCertVerifier;
@@ -38,6 +40,7 @@ pub struct Server {
     tls_acceptor: TlsAcceptor,
     listener: Arc<TcpListener>,
     handler: Handler,
+    timeout: Duration,
 }
 
 impl Server {
@@ -60,12 +63,22 @@ impl Server {
     }
 
     async fn serve_client(self, stream: TcpStream) -> Result<()> {
-        let stream = self.tls_acceptor.accept(stream).await
+        let fut_accept_request = async {
-            .context("Failed to establish TLS session")?;
+            let stream = self.tls_acceptor.accept(stream).await
-        let mut stream = BufStream::new(stream);
+                .context("Failed to establish TLS session")?;
+            let mut stream = BufStream::new(stream);
+
+            let request = receive_request(&mut stream).await
+                .context("Failed to receive request")?;
+
+            Result::<_, anyhow::Error>::Ok((request, stream))
+        };
+
+        // Use a timeout for interacting with the client
+        let fut_accept_request = timeout(self.timeout, fut_accept_request);
+        let (mut request, mut stream) = fut_accept_request.await
+            .context("Client timed out while waiting for response")??;
 
-        let mut request = receive_request(&mut stream).await
-            .context("Failed to receive request")?;
         debug!("Client requested: {}", request.uri());
 
         // Identify the client certificate from the tls stream.  This is the first
@@ -89,11 +102,18 @@ impl Server {
             })
             .context("Request handler failed")?;
 
-        send_response(response, &mut stream).await
+        // Use a timeout for sending the response
-            .context("Failed to send response")?;
+        let fut_send_and_flush = async {
+            send_response(response, &mut stream).await
+                .context("Failed to send response")?;
 
-        stream.flush().await
+            stream.flush()
-            .context("Failed to flush response data")?;
+                .await
+                .context("Failed to flush response data")
+        };
+        timeout(self.timeout, fut_send_and_flush)
+            .await
+            .context("Client timed out receiving response data")??;
 
         Ok(())
     }
@@ -103,6 +123,7 @@ pub struct Builder<A> {
     addr: A,
     cert_path: PathBuf,
     key_path: PathBuf,
+    timeout: Duration,
 }
 
 impl<A: ToSocketAddrs> Builder<A> {
@@ -111,6 +132,7 @@ impl<A: ToSocketAddrs> Builder<A> {
             addr,
             cert_path: PathBuf::from("cert/cert.pem"),
             key_path: PathBuf::from("cert/key.pem"),
+            timeout: Duration::from_secs(30),
         }
     }
 
@@ -154,6 +176,23 @@ impl<A: ToSocketAddrs> Builder<A> {
         self
     }
 
+    /// Set the timeout on incoming requests
+    ///
+    /// Note that this timeout is applied twice, once for the delivery of the request, and
+    /// once for sending the client's response.  This means that for a 1 second timeout,
+    /// the client will have 1 second to complete the TLS handshake and deliver a request
+    /// header, then your API will have as much time as it needs to handle the request,
+    /// before the client has another second to receive the response.
+    ///
+    /// If you would like a timeout for your code itself, please use
+    /// [`tokio::time::Timeout`] to implement it internally.
+    ///
+    /// The default timeout is 30 seconds.
+    pub fn set_timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = timeout;
+        self
+    }
+
     pub async fn serve<F>(self, handler: F) -> Result<()>
     where
         F: Fn(Request) -> HandlerResponse + Send + Sync + 'static,
@@ -168,6 +207,7 @@ impl<A: ToSocketAddrs> Builder<A> {
             tls_acceptor: TlsAcceptor::from(config),
             listener: Arc::new(listener),
             handler: Arc::new(handler),
+            timeout: self.timeout,
         };
 
         server.serve().await

src/types/response.rs

@@ -1,4 +1,7 @@
+use std::convert::TryInto;
+
 use anyhow::*;
+use uriparse::URIReference;
 use crate::types::{ResponseHeader, Body, Mime, Document};
 use crate::util::Cowy;
 use crate::GEMINI_MIME;
@@ -17,7 +20,7 @@ impl Response {
     }
 
     pub fn document(document: Document) -> Self {
-        Self::success(&GEMINI_MIME).with_body(document)
+        Self::success_with_body(&GEMINI_MIME, document)
     }
 
     pub fn input(prompt: impl Cowy<str>) -> Result<Self> {
@@ -35,6 +38,24 @@ impl Response {
         Self::new(header)
     }
 
+    pub fn redirect_temporary_lossy<'a>(location: impl TryInto<URIReference<'a>>) -> Self {
+        let header = ResponseHeader::redirect_temporary_lossy(location);
+        Self::new(header)
+    }
+
+    /// Create a successful response with a preconfigured body
+    ///
+    /// This is equivilent to:
+    ///
+    /// ```ignore
+    /// Response::success(mime)
+    ///     .with_body(body)
+    /// ```
+    pub fn success_with_body(mime: &Mime, body: impl Into<Body>) -> Self {
+        Self::success(mime)
+            .with_body(body)
+    }
+
     pub fn server_error(reason: impl Cowy<str>) -> Result<Self>  {
         let header = ResponseHeader::server_error(reason)?;
         Ok(Self::new(header))
@@ -45,6 +66,11 @@ impl Response {
         Self::new(header)
     }
 
+    pub fn bad_request_lossy(reason: impl Cowy<str>) -> Self {
+        let header = ResponseHeader::bad_request_lossy(reason);
+        Self::new(header)
+    }
+
     pub fn client_certificate_required() -> Self {
         let header = ResponseHeader::client_certificate_required();
         Self::new(header)

src/types/response_header.rs

@@ -1,4 +1,7 @@
+use std::convert::TryInto;
+
 use anyhow::*;
+use uriparse::URIReference;
 use crate::Mime;
 use crate::util::Cowy;
 use crate::types::{Status, Meta};
@@ -31,6 +34,18 @@ impl ResponseHeader {
         }
     }
 
+    pub fn redirect_temporary_lossy<'a>(location: impl TryInto<URIReference<'a>>) -> Self {
+        let location = match location.try_into() {
+            Ok(location) => location,
+            Err(_) => return Self::bad_request_lossy("Invalid redirect location"),
+        };
+
+        Self {
+            status: Status::REDIRECT_TEMPORARY,
+            meta: Meta::new_lossy(location.to_string()),
+        }
+    }
+
     pub fn server_error(reason: impl Cowy<str>) -> Result<Self> {
         Ok(Self {
             status: Status::PERMANENT_FAILURE,
@@ -52,6 +67,13 @@ impl ResponseHeader {
         }
     }
 
+    pub fn bad_request_lossy(reason: impl Cowy<str>) -> Self {
+        Self {
+            status: Status::BAD_REQUEST,
+            meta: Meta::new_lossy(reason),
+        }
+    }
+
     pub fn client_certificate_required() -> Self {
         Self {
             status: Status::CLIENT_CERTIFICATE_REQUIRED,

src/util.rs

@@ -5,7 +5,6 @@ use tokio::{
     fs::{self, File},
     io,
 };
-use crate::GEMINI_MIME_STR;
 use crate::types::{Response, Document, document::HeadingLevel::*};
 use itertools::Itertools;
 
@@ -20,7 +19,7 @@ pub async fn serve_file<P: AsRef<Path>>(path: P, mime: &Mime) -> Result<Response
         }
     };
 
-    Ok(Response::success(&mime).with_body(file))
+    Ok(Response::success_with_body(mime, file))
 }
 
 pub async fn serve_dir<D: AsRef<Path>, P: AsRef<Path>>(dir: D, virtual_path: &[P]) -> Result<Response> {
@@ -89,18 +88,16 @@ async fn serve_dir_listing<P: AsRef<Path>, B: AsRef<Path>>(path: P, virtual_path
 pub fn guess_mime_from_path<P: AsRef<Path>>(path: P) -> Mime {
     let path = path.as_ref();
     let extension = path.extension().and_then(|s| s.to_str());
-    let mime = match extension {
+    let extension = match extension {
-        Some(extension) => match extension {
+        Some(extension) => extension,
-            "gemini" | "gmi" => GEMINI_MIME_STR,
+        None => return mime::APPLICATION_OCTET_STREAM,
-            "txt" => "text/plain",
-            "jpeg" | "jpg" | "jpe" => "image/jpeg",
-            "png" => "image/png",
-            _ => "application/octet-stream",
-        },
-        None => "application/octet-stream",
     };
 
-    mime.parse::<Mime>().unwrap_or(mime::APPLICATION_OCTET_STREAM)
+    if let "gemini" | "gmi" = extension {
+        return crate::GEMINI_MIME.clone();
+    }
+
+    mime_guess::from_ext(extension).first_or_octet_stream()
 }
 
 /// A convenience trait alias for `AsRef<T> + Into<T::Owned>`,