Merge pull request #4 from Johann150/master

add host and port checks

README.md

@@ -27,12 +27,15 @@ openssl req -x509 -newkey rsa:4096 -keyout key.rsa -out cert.pem \
     -days 3650 -nodes -subj "/CN=example.com"
 ```
 
-3. Run the server. The command line arguments are `agate <addr:port> <content_dir> <cert_file> <key_file>`.  For example, to listen on the standard Gemini port (1965) on all interfaces:
+3. Run the server. The command line arguments are `agate <addr:port> <content_dir> <cert_file> <key_file> [<domain>]`.  For example, to listen on the standard Gemini port (1965) on all interfaces:
 
 ```
 agate 0.0.0.0:1965 path/to/content/ cert.pem key.rsa
 ```
 
+Agate will check that the port part of the requested URL matches the port specified in the 1st argument.
+If `<domain>` is specified, agate will also check that the host part of the requested URL matches this domain.
+
 When a client requests the URL `gemini://example.com/foo/bar`, Agate will respond with the file at `path/to/content/foo/bar`.  If there is a directory at that path, Agate will look for a file named `index.gmi` inside that directory.
 
 Optionally, set a log level via the `AGATE_LOG` environment variable. Logging is powered by the [env_logger crate](https://crates.io/crates/env_logger):

src/main.rs

@@ -1,7 +1,15 @@
-use async_std::{io::prelude::*, net::{TcpListener, TcpStream}, stream::StreamExt, task};
+use async_std::{
+    io::prelude::*,
+    net::{TcpListener, TcpStream},
+    stream::StreamExt,
+    task,
+};
 use async_tls::TlsAcceptor;
 use once_cell::sync::Lazy;
-use rustls::{ServerConfig, NoClientAuth, internal::pemfile::{certs, pkcs8_private_keys}};
+use rustls::{
+    internal::pemfile::{certs, pkcs8_private_keys},
+    NoClientAuth, ServerConfig,
+};
 use std::{error::Error, ffi::OsStr, fs::File, io::BufReader, marker::Unpin, sync::Arc};
 use url::Url;
 
@@ -22,18 +30,21 @@ fn main() -> Result {
     })
 }
 
-type Result<T=()> = std::result::Result<T, Box<dyn Error + Send + Sync>>;
+type Result<T = ()> = std::result::Result<T, Box<dyn Error + Send + Sync>>;
 
-static ARGS: Lazy<Args> = Lazy::new(|| args().unwrap_or_else(|| {
+static ARGS: Lazy<Args> = Lazy::new(|| {
-    eprintln!("usage: agate <addr:port> <dir> <cert> <key>");
+    args().unwrap_or_else(|| {
-    std::process::exit(1);
+        eprintln!("usage: agate <addr:port> <dir> <cert> <key> [<domain to check>]");
-}));
+        std::process::exit(1);
+    })
+});
 
 struct Args {
     sock_addr: String,
     content_dir: String,
     cert_file: String,
     key_file: String,
+    domain: Option<String>,
 }
 
 fn args() -> Option<Args> {
@@ -43,6 +54,7 @@ fn args() -> Option<Args> {
         content_dir: args.next()?,
         cert_file: args.next()?,
         key_file: args.next()?,
+        domain: args.next(),
     })
 }
 
@@ -54,16 +66,17 @@ async fn handle_request(stream: TcpStream) -> Result {
 
     let url = match parse_request(stream).await {
         Ok(url) => url,
-        Err(e) => {
+        Err((status, msg)) => {
-            respond(stream, "59", &["Invalid request."]).await?;
+            respond(stream, &status.to_string(), &[&msg]).await?;
-            return Err(e)
+            Err(msg)?
         }
     };
     if let Err(e) = send_response(url, stream).await {
         respond(stream, "51", &["Not found, sorry."]).await?;
-        return Err(e)
+        Err(e)
+    } else {
+        Ok(())
     }
-    Ok(())
 }
 
 /// TLS configuration.
@@ -80,7 +93,9 @@ fn acceptor() -> Result<TlsAcceptor> {
 }
 
 /// Return the URL requested by the client.
-async fn parse_request<R: Read + Unpin>(stream: &mut R) -> Result<Url> {
+async fn parse_request<R: Read + Unpin>(
+    stream: &mut R,
+) -> 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.
@@ -90,30 +105,48 @@ async fn parse_request<R: Read + Unpin>(stream: &mut R) -> Result<Url> {
 
     // Read until CRLF, end-of-stream, or there's no buffer space left.
     loop {
-        let bytes_read = stream.read(buf).await?;
+        let bytes_read = stream
+            .read(buf)
+            .await
+            .map_err(|_| (59, "Request ended unexpectedly"))?;
         len += bytes_read;
         if request[..len].ends_with(b"\r\n") {
             break;
         } else if bytes_read == 0 {
-            Err("Request ended unexpectedly")?
+            return Err((59, "Request ended unexpectedly"));
         }
         buf = &mut request[len..];
     }
-    let request = std::str::from_utf8(&request[..len - 2])?;
+    let request = std::str::from_utf8(&request[..len - 2]).map_err(|_| (59, "Invalid URL"))?;
 
     // Handle scheme-relative URLs.
     let url = if request.starts_with("//") {
-        Url::parse(&format!("gemini:{}", request))?
+        Url::parse(&format!("gemini:{}", request)).map_err(|_| (59, "Invalid URL"))?
     } else {
-        Url::parse(request)?
+        Url::parse(request).map_err(|_| (59, "Invalid URL"))?
     };
 
-    // Validate the URL. TODO: Check the hostname and port.
+    // Validate the URL, host and port.
     if url.scheme() != "gemini" {
-        Err("unsupported URL scheme")?
+        Err((53, "unsupported URL scheme"))
+    } else if ARGS.domain.as_ref().map_or(false, |domain| {
+        url.host().map_or(false, |host| &host.to_string() != domain)
+    }) {
+        Err((53, "proxy request refused"))
+    } else if url.port().map_or(false, |port| {
+        port != ARGS
+            .sock_addr
+            .rsplitn(2, ':')
+            .next()
+            .unwrap()
+            .parse()
+            .unwrap()
+    }) {
+        Err((59, "port did not match"))
+    } else {
+        log::info!("Got request for {:?}", url);
+        Ok(url)
     }
-    log::info!("Got request for {:?}", url);
-    Ok(url)
 }
 
 /// Send the client the file located at the requested URL.